transformer_engine/jax/flax/module.py · 87e3e56edfed81293527a798050e80a96c5fb2ee · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Fix Flax variable creation when quantizers are created directly from a recipe (#2079) · 757fd1cf
jberchtold-nvidia authored Aug 18, 2025
```
Fix flax variables when creating quantizers directly from a recipe
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
```
757fd1cf

module.py 49.6 KB