def testPrecision(self):
    ones_1d = onp.ones((2,))
    ones_2d = onp.ones((2, 2))
    ones_3d = onp.ones((2, 2, 2))
    HIGHEST = lax.Precision.HIGHEST
    jtu.assert_dot_precision(None, jnp.dot, ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.dot, precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.dot, precision=HIGHEST), ones_3d, ones_3d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.matmul, precision=HIGHEST), ones_2d, ones_2d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.vdot, precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.tensordot, axes=2, precision=HIGHEST), ones_2d, ones_2d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.tensordot, axes=(0, 0), precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.tensordot, axes=((0,), (0,)), precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.einsum, 'i,i', precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.einsum, 'ij,ij', precision=HIGHEST), ones_2d, ones_2d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.inner, precision=HIGHEST), ones_1d, ones_1d)
def testPrecision(self):
    ones_1d = onp.ones((2,))
    ones_2d = onp.ones((2, 2))
    ones_3d = onp.ones((2, 2, 2))
    HIGHEST = lax.Precision.HIGHEST
    jtu.assert_dot_precision(None, jnp.dot, ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.dot, precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.dot, precision=HIGHEST), ones_3d, ones_3d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.matmul, precision=HIGHEST), ones_2d, ones_2d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.vdot, precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.tensordot, axes=2, precision=HIGHEST), ones_2d, ones_2d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.tensordot, axes=(0, 0), precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.tensordot, axes=((0,), (0,)), precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.einsum, 'i,i', precision=HIGHEST), ones_1d, ones_1d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.einsum, 'ij,ij', precision=HIGHEST), ones_2d, ones_2d)
    jtu.assert_dot_precision(HIGHEST, partial(jnp.inner, precision=HIGHEST), ones_1d, ones_1d)