transformer_engine/jax/cpp_extensions/amax.py · dbaa02d08ef4b180f9cd61796049e82a67f636a1 · OpenDAS / TransformerEngine

jberchtold-nvidia authored Oct 09, 2025


Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

8a7ab3dd

amax.py 12.3 KB