CHANGELOG.md
CITATION.cff
CONTRIBUTING.md
LICENSE
MANIFEST.in
NOTICE
README.md
pyproject.toml
benchmarks/README.md
benchmarks/attention.md
benchmarks/attention_fa3_diff.py
benchmarks/bench_attention.py
benchmarks/bench_attention_fa3.py
benchmarks/bench_attention_forward.py
benchmarks/bench_attention_fp8_bwd.py
benchmarks/bench_context_parallel_attention.py
benchmarks/bench_conv_norm.py
benchmarks/bench_convert_hf.py
benchmarks/bench_cross_entropy.py
benchmarks/bench_dsa.py
benchmarks/bench_dsa_deepseek.py
benchmarks/bench_dsa_gather.py
benchmarks/bench_dsa_long.py
benchmarks/bench_dsa_production.py
benchmarks/bench_dsa_topk.py
benchmarks/bench_e2e_1p3b.py
benchmarks/bench_e2e_ab.py
benchmarks/bench_e2e_flame.py
benchmarks/bench_end_to_end_ab.py
benchmarks/bench_end_to_end_step.py
benchmarks/bench_expert_parallel.py
benchmarks/bench_fox.py
benchmarks/bench_fp8.py
benchmarks/bench_fp8_1p3b.py
benchmarks/bench_fp8_gemm.py
benchmarks/bench_fused_gemm.py
benchmarks/bench_fused_gemm_e2e.py
benchmarks/bench_gated_delta_rule.py
benchmarks/bench_gated_deltaproduct.py
benchmarks/bench_gdn2.py
benchmarks/bench_gdn_hd256.py
benchmarks/bench_gdn_packed.py
benchmarks/bench_generate.py
benchmarks/bench_gla.py
benchmarks/bench_hopper_features.py
benchmarks/bench_init_ab.py
benchmarks/bench_kda.py
benchmarks/bench_log_linear.py
benchmarks/bench_mamba1.py
benchmarks/bench_mamba2_dqkwg.py
benchmarks/bench_mamba2_layer.py
benchmarks/bench_mamba3.py
benchmarks/bench_mamba3_layer.py
benchmarks/bench_mamba3_mimo.py
benchmarks/bench_mixer_fast_paths.py
benchmarks/bench_mla_attention.py
benchmarks/bench_mlstm.py
benchmarks/bench_multi_gpu.py
benchmarks/bench_no_remat_1p3b.py
benchmarks/bench_pipeline.py
benchmarks/bench_ragged_dot.py
benchmarks/bench_ragged_dot_fp8.py
benchmarks/bench_remat_fp8_1p3b.py
benchmarks/bench_remat_sweep.py
benchmarks/bench_rwkv7.py
benchmarks/bench_rwkv7_bwd_ab.py
benchmarks/bench_sft_packing.py
benchmarks/bench_short_conv.py
benchmarks/bench_short_conv_e2e.py
benchmarks/bench_short_conv_l2norm.py
benchmarks/bench_short_conv_l2norm_e2e.py
benchmarks/bench_speed_of_light.py
benchmarks/bench_streaming.py
benchmarks/build_fa3_mla.sh
benchmarks/build_time_mamba1.py
benchmarks/ce_grad_accuracy_1p3b.py
benchmarks/check_mamba3_official.py
benchmarks/convert_hf.md
benchmarks/cross_entropy.md
benchmarks/diff_upstream.py
benchmarks/diff_vector_gate.py
benchmarks/dpo.md
benchmarks/dpo_trl_reference.py
benchmarks/dsa.md
benchmarks/dsa_layer_breakdown.py
benchmarks/dsa_layer_wallclock.py
benchmarks/dsa_reference_jax.py
benchmarks/dsa_reference_torch.py
benchmarks/dsa_torch_topk_ties.py
benchmarks/end_to_end.md
benchmarks/end_to_end_1p3b.md
benchmarks/expert_parallel.md
benchmarks/fox.md
benchmarks/fp8.md
benchmarks/fp8_1p3b.md
benchmarks/fsdp_step_hlo.py
benchmarks/fsdp_trace_1p3b.py
benchmarks/fused_gemm.md
benchmarks/gated_deltaproduct.md
benchmarks/gdn.md
benchmarks/gdn2_paper.md
benchmarks/gdn_warmup_error.py
benchmarks/gla.md
benchmarks/harness.py
benchmarks/init_parity.md
benchmarks/jax_tokamax_attention.py
benchmarks/kda.md
benchmarks/lean_residuals.py
benchmarks/log_linear.md
benchmarks/long_context.md
benchmarks/mamba1.md
benchmarks/mamba2.md
benchmarks/mamba3.md
benchmarks/maxtext_gdn.py
benchmarks/mla.md
benchmarks/mla_head_dims_probe.py
benchmarks/mlstm.md
benchmarks/models.md
benchmarks/multi_gpu.md
benchmarks/pipeline.md
benchmarks/profile_1p3b.md
benchmarks/profile_1p3b.py
benchmarks/profile_1p3b_flame.py
benchmarks/profile_1p3b_report.py
benchmarks/profile_fp8_passes.py
benchmarks/profile_gdn_chunk_grads.py
benchmarks/profile_gdn_stages.py
benchmarks/profile_layout_copies.py
benchmarks/profile_moe.py
benchmarks/profile_qla_stages.py
benchmarks/profile_vector_gate.py
benchmarks/ragged_dot.md
benchmarks/rwkv7.md
benchmarks/rwkv7_check.py
benchmarks/sft_packing.md
benchmarks/short_conv.md
benchmarks/speed_of_light_prototypes.py
benchmarks/speed_of_light_report.py
benchmarks/speed_of_light_step.py
benchmarks/speed_of_light_suite.py
benchmarks/speed_of_light_torch.py
benchmarks/streaming.md
benchmarks/torch_baselines.py
benchmarks/torch_cross_entropy.py
benchmarks/torch_fa3.py
benchmarks/torch_fa3_mla.py
benchmarks/torch_gla.py
benchmarks/torch_grouped_mm.py
benchmarks/torch_kda.py
benchmarks/torch_kda_diff.py
benchmarks/torch_mamba1.py
benchmarks/torch_mlstm.py
benchmarks/torch_rwkv7.py
benchmarks/torch_rwkv7_check.py
benchmarks/torch_short_conv.py
benchmarks/torch_short_conv_l2norm.py
benchmarks/torch_upstream_diff.py
benchmarks/trainer_loop_probe.py
benchmarks/xcheck_scalar_gate.py
benchmarks/xcheck_scalar_gate_torch.py
benchmarks/zero_copy_layouts.md
benchmarks/end_to_end/compare_impls.py
benchmarks/end_to_end/fw-gdn-hybrid-125m.toml
benchmarks/end_to_end/fw-gdn2-hybrid-125m.toml
benchmarks/end_to_end/fw-kda-hybrid-125m.toml
benchmarks/end_to_end/fw-mamba2-130m.toml
benchmarks/end_to_end/fw-mamba3-130m.toml
benchmarks/end_to_end/fw-transformer-125m.toml
benchmarks/end_to_end/qwen3-next-300m.toml
benchmarks/end_to_end/1p3b/gdn-1p3b.toml
benchmarks/end_to_end/1p3b/gdn-hd128-1p3b.toml
benchmarks/end_to_end/1p3b/gdn-hybrid-1p3b.toml
benchmarks/end_to_end/1p3b/gdn2-1p3b.toml
benchmarks/end_to_end/1p3b/kda-1p3b.toml
benchmarks/end_to_end/1p3b/kda-hybrid-1p3b.toml
benchmarks/end_to_end/1p3b/mamba2-1p3b.toml
benchmarks/end_to_end/1p3b/mamba3-1p3b.toml
benchmarks/end_to_end/1p3b/run_matrix.sh
benchmarks/end_to_end/1p3b/rwkv7-1p3b.toml
benchmarks/end_to_end/1p3b/transformer-1p3b.toml
benchmarks/end_to_end/1p3b/transformer-swa-1p3b.toml
benchmarks/end_to_end/1p3b/flame/check_params.py
benchmarks/end_to_end/1p3b/flame/gdn-1p3b.json
benchmarks/end_to_end/1p3b/flame/gdn-hd128-1p3b.json
benchmarks/end_to_end/1p3b/flame/job.toml
benchmarks/end_to_end/1p3b/flame/kda-1p3b.json
benchmarks/end_to_end/1p3b/flame/mamba2-1p3b.json
benchmarks/end_to_end/1p3b/flame/mamba3-1p3b.json
benchmarks/end_to_end/1p3b/flame/rwkv7-1p3b.json
benchmarks/end_to_end/1p3b/flame/setup_flame_env.sh
benchmarks/end_to_end/1p3b/flame/token_checksum.py
benchmarks/end_to_end/1p3b/flame/transformer-1p3b.json
benchmarks/end_to_end/1p3b/flame/transformer-swa-1p3b.json
benchmarks/end_to_end/1p3b/flame/fa3_shim/flash_attn/__init__.py
benchmarks/end_to_end/gdn2_paper/gdn-hybrid.toml
benchmarks/end_to_end/gdn2_paper/gdn2-hybrid.toml
benchmarks/end_to_end/gdn2_paper/kda-hybrid.toml
benchmarks/end_to_end/gdn2_paper/mamba2-hybrid.toml
benchmarks/end_to_end/gdn2_paper/mamba3-hybrid.toml
benchmarks/end_to_end/gdn2_paper/mamba3-mimo-hybrid.toml
benchmarks/end_to_end/gdn2_paper/read_paper_fig2.py
benchmarks/end_to_end/gdn2_paper/run.sh
benchmarks/end_to_end/gdn2_paper/transformer.toml
benchmarks/end_to_end/hybrids_1p3b/kimi-linear-1p3b.toml
benchmarks/end_to_end/hybrids_1p3b/qwen3-next-1p3b.toml
benchmarks/mosaic_gpu_repros/concurrent_cross_warp_reductions.py
benchmarks/mosaic_gpu_repros/inferred_swizzle_views.py
benchmarks/mosaic_gpu_repros/scoped_smem_in_branches.py
configs/gdn-hybrid-125m.toml
configs/gdn2-hybrid-125m.toml
configs/kda-hybrid-125m.toml
configs/mamba2-130m.toml
configs/mamba3-130m.toml
configs/moe-125m-ep.toml
configs/moe-125m.toml
configs/sft-chat-125m.toml
configs/tiny-stream.toml
configs/tiny-synthetic.toml
configs/transformer-125m.toml
configs/post_train/dpo-tiny.toml
docs/acknowledgements.md
docs/architectures.md
docs/kernels.md
docs/results.md
docs/training.md
docs/assets/banner-dark.png
docs/assets/banner-light.png
docs/assets/e2e-dark.png
docs/assets/e2e-light.png
docs/assets/gdn2-paper-dark.png
docs/assets/gdn2-paper-light.png
docs/assets/karlex-ai-logo.png
docs/assets/kernels-dark.png
docs/assets/kernels-light.png
docs/assets/mfu-dark.png
docs/assets/mfu-light.png
docs/assets/sweep-dark.png
docs/assets/sweep-light.png
licenses/DeepGEMM.LICENSE
licenses/DeepSeek-V3.2-Exp.LICENSE
licenses/DeepSeek-V3.LICENSE
licenses/FlashKDA.LICENSE
licenses/FlashMLA.LICENSE
licenses/FlashQLA.LICENSE
licenses/Liger-Kernel.LICENSE
licenses/MoBA.LICENSE
licenses/RWKV-LM.LICENSE
licenses/TransformerEngine.LICENSE
licenses/cutlass.LICENSE
licenses/flame.LICENSE
licenses/flash-attention.LICENSE
licenses/flash-linear-attention.LICENSE
licenses/flat.LICENSE
licenses/flax.LICENSE
licenses/jax.LICENSE
licenses/mamba.LICENSE
licenses/maxtext.LICENSE
licenses/ml-cross-entropy.LICENSE
licenses/tokamax.LICENSE
licenses/transformers.LICENSE
licenses/trl.LICENSE
licenses/xlstm.LICENSE
requirements/jax-0.11.2-cuda13.txt
src/jax_linear_attention.egg-info/PKG-INFO
src/jax_linear_attention.egg-info/SOURCES.txt
src/jax_linear_attention.egg-info/dependency_links.txt
src/jax_linear_attention.egg-info/entry_points.txt
src/jax_linear_attention.egg-info/requires.txt
src/jax_linear_attention.egg-info/top_level.txt
src/jla/__init__.py
src/jla/cli.py
src/jla/config.py
src/jla/launch_env.py
src/jla/convert/__init__.py
src/jla/convert/cli.py
src/jla/convert/hf.py
src/jla/convert/hf_config.py
src/jla/convert/hf_weights.py
src/jla/convert/safetensors_io.py
src/jla/data/__init__.py
src/jla/data/cidx.py
src/jla/data/datasets.py
src/jla/data/loader.py
src/jla/data/make_data.py
src/jla/data/packing.py
src/jla/data/pipeline.py
src/jla/data/preference.py
src/jla/data/sft.py
src/jla/data/sft_tokenize.py
src/jla/data/streaming.py
src/jla/distributed/__init__.py
src/jla/distributed/activation_sharding.py
src/jla/distributed/context_parallel.py
src/jla/distributed/context_parallel_attention.py
src/jla/distributed/expert_parallel.py
src/jla/distributed/init.py
src/jla/distributed/kernel_sharding.py
src/jla/distributed/mesh.py
src/jla/distributed/pipeline.py
src/jla/distributed/sharding.py
src/jla/eval/__init__.py
src/jla/eval/cli.py
src/jla/eval/lm_eval_adapter.py
src/jla/eval/perplexity.py
src/jla/eval/scoring.py
src/jla/inference/__init__.py
src/jla/inference/cli.py
src/jla/inference/generate.py
src/jla/inference/loading.py
src/jla/inference/rwkv7.py
src/jla/inference/sampling.py
src/jla/layers/__init__.py
src/jla/layers/attention.py
src/jla/layers/cache.py
src/jla/layers/conv.py
src/jla/layers/conv_l2norm.py
src/jla/layers/forgetting_attention.py
src/jla/layers/fp8.py
src/jla/layers/fp8_attention.py
src/jla/layers/fp8_swiglu.py
src/jla/layers/fused_mlp.py
src/jla/layers/lean.py
src/jla/layers/lightning_attention.py
src/jla/layers/linear.py
src/jla/layers/linear_attention.py
src/jla/layers/mamba1.py
src/jla/layers/mamba2_fused.py
src/jla/layers/mamba3_mimo.py
src/jla/layers/mla.py
src/jla/layers/mlp.py
src/jla/layers/moe.py
src/jla/layers/norm.py
src/jla/layers/path_attention.py
src/jla/layers/path_decode.py
src/jla/layers/rope.py
src/jla/layers/rwkv6.py
src/jla/layers/rwkv7.py
src/jla/layers/sparse_attention.py
src/jla/layers/ssm.py
src/jla/layers/ttt.py
src/jla/models/__init__.py
src/jla/models/axes.py
src/jla/models/config.py
src/jla/models/lm.py
src/jla/models/mtp.py
src/jla/models/presets.py
src/jla/models/rwkv7.py
src/jla/ops/__init__.py
src/jla/ops/chunked.py
src/jla/ops/cross_entropy.py
src/jla/ops/dsa_gather.py
src/jla/ops/log_linear.py
src/jla/ops/moba.py
src/jla/ops/moe_permute.py
src/jla/ops/path_attention.py
src/jla/ops/reference.py
src/jla/ops/rwkv6.py
src/jla/ops/rwkv7.py
src/jla/ops/sparse_attention.py
src/jla/ops/token_logprobs.py
src/jla/ops/ttt.py
src/jla/ops/mosaic_gpu/__init__.py
src/jla/ops/mosaic_gpu/_common.py
src/jla/ops/mosaic_gpu/attention.py
src/jla/ops/mosaic_gpu/attention_fp8_bwd.py
src/jla/ops/mosaic_gpu/attention_persistent.py
src/jla/ops/mosaic_gpu/dsa.py
src/jla/ops/mosaic_gpu/dsa_topk.py
src/jla/ops/mosaic_gpu/fused_gemm.py
src/jla/ops/mosaic_gpu/gated_delta_product.py
src/jla/ops/mosaic_gpu/gated_delta_rule.py
src/jla/ops/mosaic_gpu/gated_delta_rule_bwd.py
src/jla/ops/mosaic_gpu/gated_delta_rule_fused_bwd.py
src/jla/ops/mosaic_gpu/gated_delta_rule_kb.py
src/jla/ops/mosaic_gpu/gated_delta_rule_kernels.py
src/jla/ops/mosaic_gpu/gdn2.py
src/jla/ops/mosaic_gpu/gla.py
src/jla/ops/mosaic_gpu/gla_kernels.py
src/jla/ops/mosaic_gpu/kda.py
src/jla/ops/mosaic_gpu/linear_cross_entropy.py
src/jla/ops/mosaic_gpu/mamba1.py
src/jla/ops/mosaic_gpu/mamba1_bwd.py
src/jla/ops/mosaic_gpu/mamba1_bwd2.py
src/jla/ops/mosaic_gpu/mamba3.py
src/jla/ops/mosaic_gpu/mamba3_fused.py
src/jla/ops/mosaic_gpu/mamba3_mimo.py
src/jla/ops/mosaic_gpu/mamba3_mimo_kernels.py
src/jla/ops/mosaic_gpu/mlstm.py
src/jla/ops/mosaic_gpu/ragged_dot.py
src/jla/ops/mosaic_gpu/ragged_dot_fp8.py
src/jla/ops/mosaic_gpu/rwkv7.py
src/jla/ops/mosaic_gpu/rwkv7_bwd.py
src/jla/ops/mosaic_gpu/rwkv7_prep.py
src/jla/ops/mosaic_gpu/rwkv7_tiles.py
src/jla/ops/mosaic_gpu/scalar_gate.py
src/jla/ops/mosaic_gpu/scalar_gate_docs.py
src/jla/ops/mosaic_gpu/scalar_gate_fused.py
src/jla/ops/mosaic_gpu/scalar_gate_grouped.py
src/jla/ops/mosaic_gpu/scalar_gate_kb.py
src/jla/ops/mosaic_gpu/scalar_gate_kernels.py
src/jla/ops/mosaic_gpu/scalar_gate_scan.py
src/jla/ops/mosaic_gpu/short_conv.py
src/jla/ops/mosaic_gpu/short_conv_l2norm.py
src/jla/ops/mosaic_gpu/vector_gate.py
src/jla/ops/mosaic_gpu/vector_gate_bwd.py
src/jla/ops/mosaic_gpu/vector_gate_ops.py
src/jla/training/__init__.py
src/jla/training/checkpointing.py
src/jla/training/dpo.py
src/jla/training/fp8.py
src/jla/training/gpu_setup.py
src/jla/training/grad_reduce.py
src/jla/training/graph_health.py
src/jla/training/loggers.py
src/jla/training/metrics_sink.py
src/jla/training/mixed_precision.py
src/jla/training/optimizers.py
src/jla/training/performance.py
src/jla/training/runtime.py
src/jla/training/schedulers.py
src/jla/training/step.py
src/jla/training/telemetry.py
src/jla/training/trainer.py
tests/_context_parallel_attention_checks.py
tests/_sft_fakes.py
tests/_toy_model.py
tests/_vector_gate.py
tests/conftest.py
tests/test_activation_sharding.py
tests/test_bench_e2e_1p3b.py
tests/test_bench_e2e_flame_data.py
tests/test_bench_e2e_flame_report.py
tests/test_bench_end_to_end_step.py
tests/test_checkpointing.py
tests/test_chunked_ops.py
tests/test_config.py
tests/test_context_parallel.py
tests/test_context_parallel_attention.py
tests/test_context_parallel_guards.py
tests/test_conv_l2norm_plumbing.py
tests/test_convert_hf.py
tests/test_cross_entropy.py
tests/test_cudnn_fsdp_pin.py
tests/test_data_cidx.py
tests/test_data_loader.py
tests/test_data_sft.py
tests/test_data_streaming.py
tests/test_dispatch_ahead_multiprocess.py
tests/test_distributed_mesh.py
tests/test_distributed_multidevice.py
tests/test_distributed_multiprocess.py
tests/test_dpo.py
tests/test_dsa_gather.py
tests/test_eval_perplexity.py
tests/test_expert_parallel.py
tests/test_flops_presets.py
tests/test_forgetting_attention.py
tests/test_fp8_attention.py
tests/test_fp8_min_features.py
tests/test_fp8_multidevice.py
tests/test_fp8_remat_names.py
tests/test_fp8_swiglu.py
tests/test_fp8_training.py
tests/test_fp8_weight_once.py
tests/test_gated_delta_product.py
tests/test_gdn2_fast_paths.py
tests/test_gdn2_mamba3_ops.py
tests/test_generate.py
tests/test_grad_reduce.py
tests/test_grad_reduce_gpu.py
tests/test_hybrid_presets_1p3b.py
tests/test_init_parity.py
tests/test_kernel_sharding.py
tests/test_kernel_sharding_gla_mlstm_mamba1.py
tests/test_kernel_sharding_rules.py
tests/test_layers_attention.py
tests/test_layers_basic.py
tests/test_layers_conv.py
tests/test_layers_gated_deltaproduct.py
tests/test_layers_grouped_heads_precision.py
tests/test_layers_linear_attention.py
tests/test_layers_ssm.py
tests/test_lean_conv_l2norm.py
tests/test_lean_fp8.py
tests/test_lean_vjp.py
tests/test_lm_eval_adapter.py
tests/test_log_linear_layer.py
tests/test_log_linear_ops.py
tests/test_long_context.py
tests/test_mamba1.py
tests/test_mamba2_fused.py
tests/test_mamba2_fused_multidevice.py
tests/test_mamba2_in_place_zx.py
tests/test_mamba3_mimo.py
tests/test_mamba3_mimo_release.py
tests/test_mamba3_mimo_vjp.py
tests/test_mixer_fast_paths.py
tests/test_mla.py
tests/test_mlstm_layer.py
tests/test_mlstm_ops.py
tests/test_moba.py
tests/test_models_lm.py
tests/test_models_lm_extra.py
tests/test_moe.py
tests/test_moe_per_device.py
tests/test_moe_permute.py
tests/test_moe_two_gpus.py
tests/test_mosaic_gpu_attention.py
tests/test_mosaic_gpu_attention_persistent.py
tests/test_mosaic_gpu_common.py
tests/test_mosaic_gpu_dsa.py
tests/test_mosaic_gpu_dsa_topk.py
tests/test_mosaic_gpu_fused_gemm.py
tests/test_mosaic_gpu_gated_delta_rule.py
tests/test_mosaic_gpu_gdn2.py
tests/test_mosaic_gpu_gdn_hd256.py
tests/test_mosaic_gpu_gdn_split.py
tests/test_mosaic_gpu_gla.py
tests/test_mosaic_gpu_grid_limits.py
tests/test_mosaic_gpu_kda.py
tests/test_mosaic_gpu_linear_cross_entropy.py
tests/test_mosaic_gpu_linear_cross_entropy_objective.py
tests/test_mosaic_gpu_mamba1.py
tests/test_mosaic_gpu_mamba3.py
tests/test_mosaic_gpu_mamba3_fused.py
tests/test_mosaic_gpu_mamba3_mimo.py
tests/test_mosaic_gpu_mamba3_mimo_kernels.py
tests/test_mosaic_gpu_mla_attention.py
tests/test_mosaic_gpu_mlstm.py
tests/test_mosaic_gpu_ragged_dot.py
tests/test_mosaic_gpu_ragged_dot_fp8.py
tests/test_mosaic_gpu_rwkv7.py
tests/test_mosaic_gpu_scalar_gate_docs.py
tests/test_mosaic_gpu_scalar_gate_grouped.py
tests/test_mosaic_gpu_short_conv.py
tests/test_mosaic_gpu_short_conv_l2norm.py
tests/test_mosaic_gpu_vector_gate_fused.py
tests/test_mtp.py
tests/test_mtp_grad_in_forward.py
tests/test_path_attention.py
tests/test_pipeline_parallel.py
tests/test_presets_production.py
tests/test_profile_1p3b.py
tests/test_reference_ops.py
tests/test_remat_policies.py
tests/test_rope_exact_angles.py
tests/test_rwkv6_hgrn2.py
tests/test_rwkv7.py
tests/test_rwkv7_kernel_glue.py
tests/test_rwkv7_layer_bounded.py
tests/test_rwkv7_mixer_dtype.py
tests/test_rwkv7_model.py
tests/test_rwkv7_official.py
tests/test_short_conv_impl.py
tests/test_sparse_attention.py
tests/test_speed_of_light_step.py
tests/test_trainer_dispatch_ahead.py
tests/test_trainer_e2e.py
tests/test_trainer_gc_freeze.py
tests/test_trainer_resume.py
tests/test_trainer_sft.py
tests/test_trainer_stream.py
tests/test_training_gpu_setup.py
tests/test_training_graph_health.py
tests/test_training_loggers.py
tests/test_training_metrics_sink.py
tests/test_training_mixed_precision.py
tests/test_training_optimizers.py
tests/test_training_performance.py
tests/test_training_runtime.py
tests/test_training_schedulers.py
tests/test_training_step.py
tests/test_training_telemetry.py
tests/test_ttt_layer.py
tests/test_ttt_linear.py
tests/test_update_guard.py
tests/test_vector_gate_bwd_schedule.py
tests/test_vector_gate_causality.py
tests/test_vector_gate_doc_split.py
tests/test_vector_gate_gate_grads.py
tests/test_xla_tuning_1p3b.py
tests/test_zero_copy_layouts.py
tests/data/dpo_trl_reference.json