.gitignore
AGENTS.md
LICENSE
README.md
pyproject.toml
.github/workflows/ci.yml
.github/workflows/publish.yml
assets/tokenizer/chat_template.jinja
assets/tokenizer/special_tokens_map.json
assets/tokenizer/tokenizer.json
assets/tokenizer/tokenizer_config.json
benchmarks/__init__.py
benchmarks/benchmark_cema_reset.py
benchmarks/benchmark_model_paths.py
docs/README.md
docs/dev.md
docs/dtypes-and-stability.md
docs/ema-implementation.md
docs/installation.md
docs/jax-torch.md
docs/long-context-streaming.md
docs/paper-deviations.md
docs/upstream-parity-contract.md
src/megalodon_jax/__init__.py
src/megalodon_jax/_version.py
src/megalodon_jax/cache.py
src/megalodon_jax/checkpoint.py
src/megalodon_jax/config.py
src/megalodon_jax/convert.py
src/megalodon_jax/inference.py
src/megalodon_jax/model.py
src/megalodon_jax/ops.py
src/megalodon_jax/precision.py
src/megalodon_jax/types.py
src/megalodon_jax/utils.py
src/megalodon_jax.egg-info/PKG-INFO
src/megalodon_jax.egg-info/SOURCES.txt
src/megalodon_jax.egg-info/dependency_links.txt
src/megalodon_jax.egg-info/requires.txt
src/megalodon_jax.egg-info/scm_file_list.json
src/megalodon_jax.egg-info/scm_version.json
src/megalodon_jax.egg-info/top_level.txt
src/megalodon_jax/layers/__init__.py
src/megalodon_jax/layers/attention.py
src/megalodon_jax/layers/complex_ema.py
src/megalodon_jax/layers/norms.py
src/megalodon_jax/layers/rotary.py
src/megalodon_jax/layers/segments.py
src/megalodon_jax/layers/timestep_norm.py
tests/__init__.py
tests/conftest.py
tests/factories.py
tests/test_attention.py
tests/test_benchmark_infrastructure.py
tests/test_core_layers.py
tests/test_foundation.py
tests/test_inference.py
tests/test_model.py
tests/test_tooling.py
tests/test_upstream_parity.py
tests/test_workflows.py
tests/reference/__init__.py
tests/reference/cache.py
tests/reference/timestep_norm.py
tests/reference/training.py
tests/reference/upstream.py
tools/verify_modeling_correctness.py