Python np_or_jnp 예제들

예제 #1

파일: jax_util_test.py 프로젝트: 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.))

예제 #2

파일: jax_util_test.py 프로젝트: MitchellTesla/google-research

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