tests/jax/test_layer.py · 47ab4a743e2d76d7835d0c0e6f9a725e658885b1 · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Add Transformer Layer tests for pre_scale_bias and post_scale_bias (#2104) · 47ab4a74
Kshitij Lakhani authored Aug 25, 2025
```
Add Transformer Layer tests for pre_scale_bias and post_scale_bias
Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
```
47ab4a74

test_layer.py 18.9 KB