Closed copybara-service[bot] closed 8 months ago
Replace jnp.einsum with jax.lax.xeinsum where needed
Required because of https://github.com/google/jax/pull/18899
I didn't see any usage of xmap + jnp.einsum-with-'{' outside of these tests.
Replace jnp.einsum with jax.lax.xeinsum where needed
Required because of https://github.com/google/jax/pull/18899
I didn't see any usage of xmap + jnp.einsum-with-'{' outside of these tests.