Path	Lines of Code
lingvo/jax/__init__.py	1
lingvo/jax/asserts.py	256
lingvo/jax/base_input.py	140
lingvo/jax/base_layer.py	797
lingvo/jax/base_model_params.py	14
lingvo/jax/build-visibility.bzl	2
lingvo/jax/checkpoint.proto	47
lingvo/jax/checkpoint_managers.py	207
lingvo/jax/checkpoints.py	325
lingvo/jax/eval.py	554
lingvo/jax/gshard_utils.py	172
lingvo/jax/io_utils.py	8
lingvo/jax/learners.py	224
lingvo/jax/main.py	154
lingvo/jax/metric_utils.py	27
lingvo/jax/model.py	595
lingvo/jax/model_imports.py	13
lingvo/jax/model_registry.py	31
lingvo/jax/model_utils.py	120
lingvo/jax/optimizer_prefix_vectorization.py	170
lingvo/jax/optimizers.py	1172
lingvo/jax/py_utils.py	150
lingvo/jax/pytypes.py	25
lingvo/jax/schedules.py	294
lingvo/jax/summary_utils.py	253
lingvo/jax/train.py	502
lingvo/jax/train_states.py	18
lingvo/jax/trainer_lib.py	432