Python np_or_jnp Examples

Example #1

File: jax_util_test.py Project: MitchellTesla/google-research

    def test_np_or_jnp(self):
        self.assertIs(jax_util.np_or_jnp(1.), np)
        self.assertIs(jax_util.np_or_jnp(np.array(1.)), np)
        self.assertIs(jax_util.np_or_jnp(jnp.array(1.)), jnp)

        def trace_check(x):
            self.assertIs(jax_util.np_or_jnp(x), jnp)

        jax.make_jaxpr(trace_check)(np.array(1.))

Example #2

File: jax_util_test.py Project: MitchellTesla/google-research

 def trace_check(x):
     self.assertIs(jax_util.np_or_jnp(x), jnp)