text: # Copyright 2020 Google LLC
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     https://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Tests for JAX2TF converted.
  
  Specific JAX primitive conversion tests are in primitives_test."""
  
  from absl.testing import absltest
  from absl.testing import parameterized
  
  import jax
  from jax import numpy as jnp
  from jax import test_util as jtu
  from jax.config import config
  from jax.experimental import jax2tf
  from jax.experimental.jax2tf.tests import tf_test_util
  import numpy as np
  import tensorflow as tf  # type: ignore[import]
  
  config.parse_flags_with_absl()
  
  
  class Jax2TfTest(tf_test_util.JaxToTfTestCase):
  
    def test_basics(self):
      f_jax = lambda x: jnp.sin(jnp.cos(x))
      _, res_tf = self.ConvertAndCompare(f_jax, 0.7)
      self.assertIsInstance(res_tf, tf.Tensor)
  
    def test_variable_input(self):
      f_jax = lambda x: jnp.sin(jnp.cos(x))
      f_tf = jax2tf.convert(f_jax)
      v = tf.Variable(0.7)
      self.assertIsInstance(f_tf(v), tf.Tensor)
      self.assertAllClose(f_jax(0.7), f_tf(v))
  
    def test_jit(self):
      f_jax = jax.jit(lambda x: jnp.sin(jnp.cos(x)))
      self.ConvertAndCompare(f_jax, 0.7)
  
    def test_nested_jit(self):
      f_jax = jax.jit(lambda x: jnp.sin(jax.jit(jnp.cos)(x)))
      f_tf = jax2tf.convert(f_jax)
      np.testing.assert_allclose(f_jax(0.7), f_tf(0.7))
  
    def test_converts_jax_arrays(self):
      f_tf = tf.function(lambda x: x)
      self.assertEqual(f_tf(jnp.zeros([])).numpy(), 0.)
      self.assertEqual(f_tf(jnp.ones([])).numpy(), 1.)
      f_tf = tf.function(lambda x: x + x)
      self.assertEqual(f_tf(jnp.ones([])).numpy(), 2.)
  
      # Test with ShardedDeviceArray.
      n = jax.local_device_count()
      mk_sharded = lambda f: jax.pmap(lambda x: x)(f([n]))
      f_tf = tf.function(lambda x: x)
      self.assertAllClose(f_tf(mk_sharded(jnp.zeros)).numpy(),
                          np.zeros([n]))
      self.assertAllClose(f_tf(mk_sharded(jnp.ones)).numpy(),
                          np.ones([n]))
  
    @parameterized.named_parameters(jtu.cases_from_list(
      dict(testcase_name=f"_dtype={dtype.__name__}_function={with_function}",
           with_function=with_function,
           dtype=dtype)
      for with_function in [False, True]
      for dtype in [np.int64, np.float64]))
    def test_converts_64bit(self, dtype=np.int64, with_function=False):
      if jtu.device_under_test() == "tpu":
        # TODO(necula): fix and re-enable this test.
        raise self.skipTest("Test is flaky on TPU")
      big_const = np.full((5,), 2 ** 33, dtype=dtype)
      self.ConvertAndCompare(jnp.sin, big_const)
      f_conv = jax2tf.convert(jnp.sin)
      if with_function:
        f_conv = tf.function(f_conv)
      # We check also when we pass tf.Variable or tf.Tensor into the
      # converted function
      self.assertAllClose(jnp.sin(big_const),
                          f_conv(tf.Variable(big_const)))
      self.assertAllClose(jnp.sin(big_const),
                          f_conv(tf.constant(big_const)))
  
    def test_function(self):
      f_jax = jax.jit(lambda x: jnp.sin(jnp.cos(x)))
      self.ConvertAndCompare(f_jax, 0.7, with_function=True)
  
    def test_gradients_disabled(self):
      f = jax2tf.convert(jnp.tan)
      x = tf.ones([])
      with tf.GradientTape() as tape:
        tape.watch(x)
        y = f(x)
      with self.assertRaisesRegex(ValueError,
                                  'jax2tf currently does not support gradients'):
        tape.gradient(y, x)
  
    def test_convert_argument_non_callable_error(self):
      with self.assertRaisesRegex(TypeError, "Expected a callable value"):
        jax2tf.convert(5.)
  
    def test_convert_argument_non_tensor_error(self):
      with self.assertRaisesRegex(TypeError,
                                  "Argument.*should be NumPy array"):
        jax2tf.convert(lambda x: x)(lambda y: y)
  
    def test_argument_eager_tensor(self):
      x = jax2tf.convert(jnp.sin)(1.)
      jax2tf.convert(jnp.cos)(x)  # No error
  
    def test_checkpoint_wrapper_types(self):
      m = tf.Module()
      m.a = [tf.Module(), tf.Module()]
      m.b = (tf.Module(), tf.Module())
      m.c = {'a': tf.Module(), 'b': tf.Module()}
      self.assertNotEqual(type(m.a), list)
      self.assertNotEqual(type(m.b), tuple)
      self.assertNotEqual(type(m.c), dict)
      self.assertLen(jax.tree_leaves(m.a), 2)
      self.assertLen(jax.tree_leaves(m.b), 2)
      self.assertLen(jax.tree_leaves(m.c), 2)
  
  
  if __name__ == "__main__":
    absltest.main()
  