docs
examples
kfac_jax
