├── .gitignore ├── LICENSE ├── README.md ├── benchmarks ├── attn.py └── varlen.py ├── load_model_with_dolomite_demo.py ├── setup.py ├── stickbreaking_attention ├── __init__.py ├── sb_attn │ ├── __init__.py │ ├── sb_bwd.py │ └── sb_fwd.py ├── sb_ref.py ├── sb_varlen │ ├── __init__.py │ ├── sb_varlen_bwd.py │ ├── sb_varlen_fwd.py │ └── softplus.py └── utils.py └── tests ├── __init__.py ├── test_attn.py └── test_varlen.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/README.md -------------------------------------------------------------------------------- /benchmarks/attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/benchmarks/attn.py -------------------------------------------------------------------------------- /benchmarks/varlen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/benchmarks/varlen.py -------------------------------------------------------------------------------- /load_model_with_dolomite_demo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/load_model_with_dolomite_demo.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/setup.py -------------------------------------------------------------------------------- /stickbreaking_attention/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/__init__.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_attn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_attn/__init__.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_attn/sb_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_attn/sb_bwd.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_attn/sb_fwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_attn/sb_fwd.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_ref.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_ref.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_varlen/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_varlen/__init__.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_varlen/sb_varlen_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_varlen/sb_varlen_bwd.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_varlen/sb_varlen_fwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_varlen/sb_varlen_fwd.py -------------------------------------------------------------------------------- /stickbreaking_attention/sb_varlen/softplus.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/sb_varlen/softplus.py -------------------------------------------------------------------------------- /stickbreaking_attention/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/stickbreaking_attention/utils.py -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/tests/test_attn.py -------------------------------------------------------------------------------- /tests/test_varlen.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shawntan/stickbreaking-attention/HEAD/tests/test_varlen.py --------------------------------------------------------------------------------