LICENSE
README.md
pyproject.toml
src/haliax/__init__.py
src/haliax/axis.py
src/haliax/core.py
src/haliax/debug.py
src/haliax/hof.py
src/haliax/jax_utils.py
src/haliax/ops.py
src/haliax/partitioning.py
src/haliax/random.py
src/haliax/specialized_fns.py
src/haliax/tree_util.py
src/haliax/types.py
src/haliax/util.py
src/haliax/wrap.py
src/haliax.egg-info/PKG-INFO
src/haliax.egg-info/SOURCES.txt
src/haliax.egg-info/dependency_links.txt
src/haliax.egg-info/requires.txt
src/haliax.egg-info/top_level.txt
src/haliax/_src/__init__.py
src/haliax/_src/util.py
src/haliax/nn/__init__.py
src/haliax/nn/attention.py
src/haliax/nn/dropout.py
src/haliax/nn/embedding.py
src/haliax/nn/linear.py
src/haliax/nn/normalization.py
src/haliax/nn/scan.py
tests/test_attention.py
tests/test_axis.py
tests/test_debug.py
tests/test_hof.py
tests/test_nn.py
tests/test_ops.py
tests/test_partitioning.py
tests/test_random.py
tests/test_specialized_fns.py
tests/test_tree_util.py
tests/test_utils.py