jax

ccb8d459 - Uses jnp.square instead of power. (#3036)

Uses jnp.square instead of power. (#3036) * Uses multiplication instead of power. * Uses jnp.square instead of mul and adds check if jnp.square is implemented by mul.