tests/jax/test_layer.py · 82e5b4d2fa0ec5f6a797253d5de4a85b42b470ae · OpenDAS / TransformerEngine

[JAX] Fixed the shape miss-matching issue in MLP. (#859) · 82e5b4d2

Ming-Xu Huang authored May 22, 2024



* Fixed the shape mismatching issue in MLP.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Add a corresponding test
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

82e5b4d2

test_layer.py 15.8 KB

Replace test_layer.py