.gitignore
.pre-commit-config.yaml
CONTRIBUTING.md
LICENSE
README.md
mkdocs.yml
pyproject.toml
.github/workflows/ci.yml
.github/workflows/docs.yml
docs/index.md
docs/requirements.txt
docs/triton_call.md
examples/JAX_+_Triton_Flash_Attention.ipynb
examples/add.py
examples/fused_attention.py
examples/matmul.py
examples/softmax.py
examples/fusion/benchmark_matmul.py
examples/fusion/nn.py
jax_triton/__init__.py
jax_triton/triton_lib.py
jax_triton/utils.py
jax_triton/version.py
jax_triton.egg-info/PKG-INFO
jax_triton.egg-info/SOURCES.txt
jax_triton.egg-info/dependency_links.txt
jax_triton.egg-info/requires.txt
jax_triton.egg-info/top_level.txt
jax_triton/experimental/fusion/__init__.py
jax_triton/experimental/fusion/fusion.py
jax_triton/experimental/fusion/jaxpr_rewriter.py
jax_triton/experimental/fusion/lowering.py
tests/cluster_test.py
tests/triton_call_test.py
tests/triton_test.py