Replace manual PrecisionLike definition with jax.lax.PrecisionLike and remove unused lax import

google-research / e3x

E3x is a JAX library for constructing efficient E(3)-equivariant deep learning architectures built on top of Flax.

https://e3x.readthedocs.io/stable/

Apache License 2.0

87 stars 5 forks source link

Closed copybara-service[bot] closed 9 months ago

copybara-service[bot] commented 9 months ago

Replace manual PrecisionLike definition with jax.lax.PrecisionLike and remove unused lax import