jax
d2a39bc6 - Updated the layer norm implementation in Mosaic GPU tests

Commit

1 year ago

Updated the layer norm implementation in Mosaic GPU tests jnp.var now needs lax.gt_p, which we don't currently support. PiperOrigin-RevId: 639011383

References

Author

superbobry

Committer

a-googler

Parents