Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
14 lines
336 B
Python
14 lines
336 B
Python
"""Readable PyTorch implementations used as correctness references."""
|
|
|
|
from .chunkwise import naive_chunk_kda
|
|
from .gate import kda_gate_naive, kda_gate_reference
|
|
from .recurrent import naive_kda, naive_kda_fwd
|
|
|
|
__all__ = [
|
|
"kda_gate_naive",
|
|
"kda_gate_reference",
|
|
"naive_chunk_kda",
|
|
"naive_kda",
|
|
"naive_kda_fwd",
|
|
]
|