examples/jax/mnist/test_single_gpu_mnist.py · d99142a0177a2462cbda07a31aaa8e68b4e85461 · OpenDAS / TransformerEngine

"...git@developer.sourcefind.cn:kecinstone/2024-pra-vllm.git" did not exist on "9f88db35da012544a9cd9450c68b8df2e6509b92"

Find file Blame History Permalink

[JAX] Allow multi-dims for dgamma and dbeta in LN descriptor. (#780) · aaf93548

Ming-Xu Huang authored Apr 19, 2024



* Allow multi-dims for dgamma and dbeta in LN descriptor.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix the jit error in examples/jax
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>

aaf93548

test_single_gpu_mnist.py 10 KB

Replace test_single_gpu_mnist.py