LICENSE
MANIFEST.in
README.md
pyproject.toml
setup.cfg
setup.py
xax/__init__.py
xax/py.typed
xax/requirements-dev.txt
xax/requirements.txt
xax.egg-info/PKG-INFO
xax.egg-info/SOURCES.txt
xax.egg-info/dependency_links.txt
xax.egg-info/entry_points.txt
xax.egg-info/requires.txt
xax.egg-info/top_level.txt
xax/cli/__init__.py
xax/cli/edit_config.py
xax/core/__init__.py
xax/core/conf.py
xax/core/state.py
xax/nn/__init__.py
xax/nn/attention.py
xax/nn/distributions.py
xax/nn/embeddings.py
xax/nn/functions.py
xax/nn/geom.py
xax/nn/losses.py
xax/nn/metrics.py
xax/nn/parallel.py
xax/nn/ssm.py
xax/task/__init__.py
xax/task/base.py
xax/task/logger.py
xax/task/script.py
xax/task/task.py
xax/task/launchers/__init__.py
xax/task/launchers/base.py
xax/task/launchers/cli.py
xax/task/launchers/single_process.py
xax/task/loggers/__init__.py
xax/task/loggers/callback.py
xax/task/loggers/json.py
xax/task/loggers/state.py
xax/task/loggers/stdout.py
xax/task/loggers/tensorboard.py
xax/task/mixins/__init__.py
xax/task/mixins/artifacts.py
xax/task/mixins/checkpointing.py
xax/task/mixins/compile.py
xax/task/mixins/cpu_stats.py
xax/task/mixins/data_loader.py
xax/task/mixins/gpu_stats.py
xax/task/mixins/logger.py
xax/task/mixins/process.py
xax/task/mixins/runnable.py
xax/task/mixins/step_wrapper.py
xax/task/mixins/train.py
xax/utils/__init__.py
xax/utils/debugging.py
xax/utils/experiments.py
xax/utils/jax.py
xax/utils/jaxpr.py
xax/utils/logging.py
xax/utils/numpy.py
xax/utils/profile.py
xax/utils/pytree.py
xax/utils/tensorboard.py
xax/utils/text.py
xax/utils/data/__init__.py
xax/utils/data/collate.py
xax/utils/types/__init__.py
xax/utils/types/frozen_dict.py
xax/utils/types/hashable_array.py