Functions
build_lr_schedule(base_lr, schedule_config)
build_lr_schedule
build_optax_optimizer(config)
build_optax_optimizer
build_optimizer(value_and_grad_func, opt_config)
build_optimizer
Classes
OptaxWrapper(value_and_grad_func, ...[, ...])
OptaxWrapper
Wrapper class for Optax optimizers to have the same interface as KFAC.