LICENSE
README.md
pyproject.toml
setup.py
FJFormer.egg-info/PKG-INFO
FJFormer.egg-info/SOURCES.txt
FJFormer.egg-info/dependency_links.txt
FJFormer.egg-info/requires.txt
FJFormer.egg-info/top_level.txt
fjformer/__init__.py
fjformer/utils.py
fjformer/attention/__init__.py
fjformer/attention/efficient_attention.py
fjformer/attention/flash_attention.py
fjformer/attention/flash_attention_0.py
fjformer/attention/flash_attention_gpu.py
fjformer/attention/flash_attention_tpu.py
fjformer/attention/jax_flash_attn_gpu.py
fjformer/attention/jax_flash_attn_tpu.py
fjformer/bits/__init__.py
fjformer/bits/bits.py
fjformer/bits/calibration.py
fjformer/bits/config.py
fjformer/bits/int_numerics.py
fjformer/bits/no_numerics.py
fjformer/bits/numerics.py
fjformer/bits/q_dot_general.py
fjformer/bits/q_flax.py
fjformer/bits/stochastic_rounding.py
fjformer/checkpoint/__init__.py
fjformer/checkpoint/_load.py
fjformer/checkpoint/streamer.py
fjformer/datasets/__init__.py
fjformer/datasets/datasets.py
fjformer/func/__init__.py
fjformer/func/_func.py
fjformer/func/loss_func.py
fjformer/gpu_pallas/__init__.py
fjformer/gpu_pallas/attention.py
fjformer/gpu_pallas/layer_norm.py
fjformer/gpu_pallas/rms_norm.py
fjformer/gpu_pallas/softmax.py
fjformer/monitor/__init__.py
fjformer/monitor/tracker.py
fjformer/optimizers/__init__.py
fjformer/optimizers/adafactor.py
fjformer/optimizers/adamw.py
fjformer/optimizers/lion.py
fjformer/optimizers/optimizer_utils.py
fjformer/partition_utils/__init__.py
fjformer/partition_utils/mesh_utils.py
fjformer/partition_utils/t5x_partitioning.py