Commits · 87e3e56edfed81293527a798050e80a96c5fb2ee · OpenDAS / TransformerEngine

30 Jul, 2025 1 commit

[JAX] TE GEMM checkpointing policies (#2003) · 858755c0

jberchtold-nvidia authored Jul 30, 2025



* TE primitive checkpointing policies
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

* Remove batched gemm policy
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

858755c0