Tree Attention Torch

An implementation of Tree-Attention in PyTorch because it's in JAX for some reason

Usage

python3 model.py

MIT

Implement flash attention from the native official repo, I couldn't because the docs are nowhere to be found and understood

Name		Name	Last commit message	Last commit date
Latest commit History 5 Commits
.github		.github
scripts		scripts
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
Dockerfile		Dockerfile
LICENSE		LICENSE
README.md		README.md
agorabanner.png		agorabanner.png
model.py		model.py
pyproject.toml		pyproject.toml
requirements.txt		requirements.txt