flax
4e300d4c - Avoid calls to jnp.shape for non-array inputs.

Commit

292 days ago

Avoid calls to jnp.shape for non-array inputs. These are deprecated in https://github.com/jax-ml/jax/pull/26641

References

Author

jakevdp

Parents