Official JAX implementation of neural isometries - taming transformations for equivariant ML - View it on GitHub