# Copyright 2020 The JAX Authors.
#
# 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 call_tf."""
from functools import partial
import os
from typing import Callable
import unittest

from absl import logging
from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import config
from jax import dtypes
from jax import lax
from jax import numpy as jnp
from jax._src import test_util as jtu
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import hlo
from jax.experimental import jax2tf
from jax.experimental.jax2tf.tests import tf_test_util
import numpy as np

try:
  import tensorflow as tf  # type: ignore[import]
except ImportError:
  tf = None

config.parse_flags_with_absl()


def _maybe_jit(with_jit: bool, func: Callable) -> Callable:
  if with_jit:
    return jax.jit(func)
  else:
    return func

def _maybe_tf_jit(with_jit: bool, func: Callable) -> Callable:
  if with_jit:
    return tf.function(func, autograph=False, jit_compile=True)
  else:
    return func

def _named_test(**kwargs):
  return dict(kwargs,
              testcase_name = "_".join([f"{k}={kwargs[k]}" for k in sorted(kwargs.keys())]))

_parameterized_jit = parameterized.named_parameters(
    _named_test(with_jit=with_jit)
    for with_jit in [True, False])

_call_tf_non_compilable_error = "Error compiling TensorFlow function"
_call_tf_dynamic_shape_error = "call_tf cannot call functions whose output has dynamic shape"

class CallTfTest(tf_test_util.JaxToTfTestCase):

  def setUp(self):
    if tf is None:
      raise unittest.SkipTest("Test requires tensorflow")
    # TODO(b/171320191): this line works around a missing context initialization
    # bug in TensorFlow.
    _ = tf.add(1, 1)
    super().setUp()

  @_parameterized_jit
  def test_eval_scalar_arg(self, with_jit=True):
    def f_tf(x):
      return tf.math.sin(x)
    x = 3.
    res = _maybe_jit(with_jit, jax2tf.call_tf(f_tf))(x)
    self.assertAllClose(jnp.sin(x), res)

  @_parameterized_jit
  def test_eval_scalar_res(self, with_jit=True):
    x = 3.
    res = _maybe_jit(with_jit, jax2tf.call_tf(lambda x: 4.))(x)
    self.assertAllClose(4., res, check_dtypes=False)

  @_parameterized_jit
  def test_eval_numpy_arg(self, with_jit=True):
    x = np.ones((2, 3), dtype=np.float32)
    res = _maybe_jit(with_jit, jax2tf.call_tf(tf.math.sin))(x)
    self.assertAllClose(jnp.sin(x), res)

  @_parameterized_jit
  def test_eval_numpy_res(self, with_jit=False):
    x = np.ones((2, 3))
    res = _maybe_jit(with_jit, jax2tf.call_tf(lambda _: x))(x)
    self.assertAllClose(x, res)

  @_parameterized_jit
  def test_eval_devicearray_arg(self, with_jit=False):
    x = jnp.ones((2, 3), dtype=np.float32)
    res = _maybe_jit(with_jit, jax2tf.call_tf(tf.math.sin))(x)
    self.assertAllClose(jnp.sin(x), res)

    x = jnp.array(3.0, dtype=jnp.bfloat16)
    res = jax2tf.call_tf(lambda x: x)(x)
    self.assertAllClose(x, res)
    # bfloat16 scalar will create a copy.
    with self.assertRaises(AssertionError):
      self.assertTrue(np.shares_memory(x, res))

  @_parameterized_jit
  def test_eval_pytree(self, with_jit=True):

    def fun_tf(x: dict, y: tuple) -> tuple:
      return (x["first"] * x["second"], y[0] + y[1])

    x = dict(first=np.float32(3.), second=np.float32(4.))
    y = (np.float64(5.), np.float64(6.))
    fun_jax = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))
    res = fun_jax(x, y)
    self.assertAllClose((np.float32(12.), np.float64(11.)), res)

  def test_result_tuple(self):
    x1 = np.ones(3, dtype=np.int32)
    x2 = np.ones(5, dtype=np.float32)
    def fun_tf():
      return tf.tuple([x1, x2])

    fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
    res = fun_jax()
    self.assertAllClose(res, (x1, x2))

  def test_error_non_compilable_strings(self):
    # Check that in op-by-op we call a function in eager mode.
    def f_tf_non_compilable(x):
      return tf.strings.length(tf.strings.format("Hello {}!", [x]))

    f_jax = jax2tf.call_tf(f_tf_non_compilable)
    x = np.float32(0.7)
    self.assertAllClose(f_tf_non_compilable(x).numpy(), f_jax(x))
    with self.assertRaisesRegex(ValueError,
                                _call_tf_non_compilable_error):
      jax.jit(f_jax)(x)

    with self.assertRaisesRegex(ValueError,
                                _call_tf_non_compilable_error):
      lax.cond(True, lambda x: f_jax(x), lambda x: f_jax(x), x)

  def test_error_non_compilable_dynamic_shape(self):
    # Check that in op-by-op we call a function in eager mode.
    def f_tf_non_compilable(x):
      return tf.cond(x[0], lambda: x[1:], lambda: x)

    f_jax = jax2tf.call_tf(f_tf_non_compilable)
    x = np.array([True, False], dtype=np.bool_)
    self.assertAllClose(f_tf_non_compilable(x), f_jax(x))  # Works in eager mode

    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
      jax.jit(f_jax)(x)

  def test_error_bad_result_tensorarray(self):
    # Call a function that returns a tf.TensorArray. This should be detected
    # early on. If we don't the function is actually compilable but returns
    # a tuple instead of a single result.
    def fun_tf():
      ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
      ta = ta.unstack([0, 1, 2, 3, 4])
      return ta

    with self.assertRaisesRegex(ValueError,
                                "The called TF function returns a result that is not convertible to JAX"):
      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
      fun_jax()

  def test_error_bad_result_string(self):
    def fun_tf():
      return tf.constant("foo")

    # Now under jit, should fail because the function is not compilable
    with self.assertRaisesRegex(ValueError,
                                "The called TF function returns a result that is not convertible to JAX"):
      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
      fun_jax()

  @_parameterized_jit
  def test_control_flow(self, with_jit=True):

    def times_5_tf(x):
      # Multiply x * 5 using a loop
      c = lambda i, acc: tf.less(i, 5)
      b = lambda i, acc: (tf.add(i, 1), tf.add(acc, x))
      _, acc = tf.while_loop(c, b, [tf.constant(0), tf.constant(0.)])
      return acc

    def fun_jax(x):
      # Calls times_5_tf 3 times in a loop
      def body(_, acc):
        return jax2tf.call_tf(times_5_tf)(acc)

      return lax.fori_loop(0, 3, body, x)

    x = np.float32(3.)
    res = _maybe_jit(with_jit, fun_jax)(x)
    self.assertAllClose(np.float32(x * 5 * 5 * 5), res)

  @parameterized.named_parameters(
      dict(
          testcase_name=f"_{dtype.__name__}{'_jit' if with_jit else ''}",
          dtype=dtype,
          with_jit=with_jit)
      for dtype in set(jtu.dtypes.all) - {np.bool_}
      for with_jit in [True, False])
  def test_dtypes(self, dtype=np.int32, with_jit=True):

    def fun_tf(x):
      # AddV2 supports more types
      return tf.raw_ops.AddV2(x=x, y=tf.constant(3, dtype=dtype))

    def fun_jax(x):
      return jax2tf.call_tf(fun_tf)(x) + x

    x = np.ones((3,), dtype=dtype)
    res = _maybe_jit(with_jit, fun_jax)(x)
    self.assertAllClose(dtype(2 * x + 3), res)

  @_parameterized_jit
  def test_bool(self, with_jit=False):

    def fun_tf(x, y):
      return tf.math.logical_and(x, y)

    x = np.array([True, False, True, False], dtype=np.bool_)
    y = np.array([True, True, False, False], dtype=np.bool_)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x, y)
    self.assertAllClose(
        np.array([True, False, False, False], dtype=np.bool_), res)

  @_parameterized_jit
  def test_x64_input(self, with_jit=True):
    def f_tf(x):
      return tf.math.sin(x)

    x = 5.  # TF interprets this as f64
    res_call_tf = _maybe_jit(with_jit, jax2tf.call_tf(f_tf))(x)
    res_jax = jnp.sin(x)
    self.assertAllClose(res_call_tf, res_jax)

  @_parameterized_jit
  def test_x64_output(self, with_jit=True):
    def f_tf(x):
      return (tf.constant(3., tf.float64), x)

    x = np.float32(5.)
    res_call_tf = _maybe_jit(with_jit, jax2tf.call_tf(f_tf))(x)
    res_jax = (3., x)
    self.assertAllClose(res_call_tf, res_jax)

    res_call_tf_jit = jax.jit(jax2tf.call_tf(f_tf))(x)
    self.assertAllClose(res_call_tf_jit, res_jax)

  @_parameterized_jit
  def test_with_var_read(self, with_jit=True):
    # The variable is placed on the default TF device.
    outer_var_array = np.array([3., 4.], dtype=np.float32)
    outer_var = tf.Variable(outer_var_array)

    def fun_tf(x):
      return x * outer_var + 1.

    x = np.array([2., 5.,], dtype=np.float32)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * outer_var_array + 1., res, check_dtypes=False)

  @_parameterized_jit
  def test_with_var_read_x64(self, with_jit=True):
    outer_var_array = np.array([3., 4.], dtype=np.float64)
    outer_var = tf.Variable(outer_var_array)

    def fun_tf(x):
      return x * tf.cast(outer_var, x.dtype) + 1.

    x = np.array([2., 5.,], dtype=np.float32)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * outer_var_array + 1., res, check_dtypes=False)

  def test_with_var_different_shape(self):
    # See https://github.com/google/jax/issues/6050
    v = tf.Variable((4., 2.), dtype=tf.float32)

    def tf_func(x):
      return v + x
    x = np.float32(123.)
    tf_out = tf_func(x)

    jax_func = jax.jit(jax2tf.call_tf(tf_func))
    jax_out = jax_func(x)

    self.assertAllClose(tf_out, jax_out, check_dtypes=False)

  @_parameterized_jit
  def test_with_var_write_error(self, with_jit=True):
    if with_jit:
      raise unittest.SkipTest("variable writes not yet working")
    outer_var = tf.Variable(3., dtype=np.float32)

    def fun_tf(x):
      outer_var.assign(tf.constant(4.))
      return x * outer_var + 1.

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * 4. + 1, res, check_dtypes=False)

  @_parameterized_jit
  def test_with_tensor_capture(self, with_jit=True):
    outer_tensor = tf.constant(3., dtype=np.float32)

    def fun_tf(x):
      return x * outer_tensor + 1.

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * 3. + 1., res, check_dtypes=False)

  @_parameterized_jit
  def test_with_tensor_capture_x64(self, with_jit=True):
    outer_tensor = tf.constant(3., dtype=np.float64)

    def fun_tf(x):
      return x * tf.cast(outer_tensor * 3.14, tf.float32) + 1.

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * 3. * 3.14 + 1., res, check_dtypes=False)

  @_parameterized_jit
  def test_with_value_capture(self, with_jit=True):
    outer_val = np.array(3., dtype=np.float32)

    def fun_tf(x):
      return x * outer_val + 1.

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose(x * 3. + 1., res, check_dtypes=False)

  @_parameterized_jit
  def test_with_multiple_capture(self, with_jit=True):
    if jtu.device_under_test() == "gpu":
      raise unittest.SkipTest("Test fails on GPU")
    v2 = tf.Variable(2., dtype=np.float32)
    v3 = tf.Variable(3., dtype=np.float32)
    t4 = tf.constant(4., dtype=np.float32)
    t5 = tf.constant(5., dtype=np.float32)

    def fun_tf(x):
      return (x * v3 + t4 + v2) * v3 + t5

    x = np.float32(2.)
    res = _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)
    self.assertAllClose((x * 3. + 4. + 2.) * 3. + 5., res, check_dtypes=False)

  @_parameterized_jit
  def test_grad(self, with_jit=False):
    x = np.float32(3.)
    res = _maybe_jit(with_jit, jax.grad(jax2tf.call_tf(tf.math.sin)))(x)
    self.assertAllClose(np.cos(x), res)

  @_parameterized_jit
  def test_grad_pytree(self, with_jit=False):

    def fun_tf(x: dict, y: tuple) -> tuple:
      return x["first"] * x["second"] + 3. * y[0] + 4. * y[1]

    x = dict(first=np.float32(3.), second=np.float32(4.))
    y = (np.float32(5.), np.float32(6.))
    grad_x = _maybe_jit(with_jit, jax.grad(jax2tf.call_tf(fun_tf)))(x, y)
    self.assertAllClose(
        dict(first=np.float32(4.), second=np.float32(3.)), grad_x)

  def test_grad_nested(self):
    # We embed the call_tf function in a larger function whose gradient we take
    # It is relevant here that the cotangents flowing through the call_tf
    # function are not scalars.

    b = np.array([[11., 12., 13.], [21., 22., 23.]], dtype=np.float32)  # [2, 3]
    c = np.array([[31., 32.], [41., 42.], [51., 52.], [61., 62.]], dtype=np.float32)  # [4, 2]
    x_dict = dict(b=b, c=c)  # b:[2, 3], c=[4, 2]
    # res: dict(r:[4, 3], s:[4, 2])
    def f_tf(x_dict):
      return dict(r=tf.matmul(x_dict["c"], x_dict["b"]), s=7. * x_dict["c"])

    @jax.jit  # To recognize it in jaxpr
    def f_jax(x_dict):
      return dict(r=jnp.matmul(x_dict["c"], x_dict["b"]), s=7. * x_dict["c"])

    def loss(functional, x_dict):
      prediction = functional(x_dict)  # r:[4, 3], s:[4, 2]
      weights = np.array([1., 2., 3., 4.], dtype=np.float32)  # [4]
      weighted_pred = jnp.matmul(weights, prediction["r"])  # [3]
      return jnp.sum(weighted_pred) + 4. * jnp.sum(prediction["s"])

    g_fun_with_tf = jax.grad(partial(loss, jax2tf.call_tf(f_tf)))
    g_fun_with_jax = jax.grad(partial(loss, f_jax))

    g_tf = g_fun_with_tf(x_dict)
    g_jax = g_fun_with_jax(x_dict)
    self.assertAllClose(g_jax, g_tf)

  def test_grad_int_argument(self):
    # Similar to https://github.com/google/jax/issues/6975
    # state is a pytree that contains an integer and a boolean.
    # The function returns an integer and a boolean.
    def f(param, state, x):
      return param * x, state

    param = np.array([0.7, 0.9], dtype=np.float32)
    state = dict(array=np.float32(1.), counter=7, truth=True)
    x = np.float32(3.)

    # tf.function is important, without it the bug does not appear
    f_call_tf = jax2tf.call_tf(f)
    g_call_tf = jax.grad(lambda *args: jnp.sum(f_call_tf(*args)[0]))(param, state, x)
    g = jax.grad(lambda *args: jnp.sum(f(*args)[0]))(param, state, x)
    self.assertAllClose(g_call_tf, g)

  def test_grad_int_argument_unused(self):
    batch_size = 5
    inputs = np.ones((batch_size, 3), dtype=np.float32)
    rng = np.array([1, 2], dtype=np.uint32)
    params = np.float32(.5)

    # rng is integer, unused
    def jax_model(params, rng, inputs):
      return jnp.ones([batch_size, 2], dtype=jnp.float32)

    tf_model = jax2tf.convert(jax_model, with_gradient=True)

    def _loss_fn(inference_fn, params, rng, inputs):
      prediction = inference_fn(params, rng, inputs)
      return jnp.mean(prediction)

    jax_loss_fn = partial(_loss_fn, jax_model)
    jax_grad = jax.grad(jax_loss_fn)(params, rng, inputs)

    paramsv = tf.Variable(params)
    with tf.GradientTape() as tape:
      tf_prediction = tf_model(paramsv, rng, inputs)
      tf_loss = tf.reduce_mean(tf_prediction)

      tf_grad = tape.gradient(tf_loss, paramsv)
    self.assertAllClose(jax_grad, tf_grad.numpy())

    call_tf_loss_fn = partial(_loss_fn, jax2tf.call_tf(tf_model))
    call_tf_grad = jax.grad(call_tf_loss_fn)(params, rng, inputs)
    self.assertAllClose(jax_grad, call_tf_grad)

  def test_grad_with_float0_result(self):
    # Gradient over integer-argument functions, with float0 result
    def f_jax(x, y):  # x is an int, y is a float; res is a (int, float)
      return (2 * x, 2 * x + y * y)
    def f_tf(x, y):
      # TF needs explicit casts
      return (2 * x, tf.cast(2 * x, dtype=y.dtype) + y * y)

    def wrapper(functional, x, y):  # x: i32
      return jnp.sum(2. * functional(3 * x, 4. * y)[1])

    grad_g = jax.grad(partial(wrapper, f_jax),
                      allow_int=True, argnums=(0, 1))
    grad_g_call_tf = jax.grad(partial(wrapper, jax2tf.call_tf(f_tf)),
                              allow_int=True, argnums=(0, 1))

    x = np.int32(2)
    y = np.float32(3.)
    g_jax = grad_g(x, y)
    g_call_tf = grad_g_call_tf(x, y)
    self.assertEqual(g_jax[0].dtype, dtypes.float0)
    self.assertEqual(g_call_tf[0].dtype, dtypes.float0)
    self.assertAllClose(g_jax[1], g_call_tf[1])

  @_parameterized_jit
  def test_grad_custom(self, with_jit=False):

    @tf.custom_gradient
    def func_square_tf(x):
      # Like x ** 2, but with custom grad 3. * x
      def grad(dy, variables=None):
        # dy, = dys
        return 3. * x * dy,

      return x * x, grad

    x = np.float32(4.)
    grad_x = _maybe_jit(with_jit, jax.grad(jax2tf.call_tf(func_square_tf)))(x)
    self.assertAllClose(np.float32(3.) * x, grad_x)

  @parameterized.named_parameters(
      dict(
          testcase_name=f"_{degree=}{'_jit' if with_jit else ''}",
          degree=degree,
          with_jit=with_jit)
      for degree in [1, 2, 3, 4]
      for with_jit in [True, False])
  def test_higher_order_grad(self, degree=2, with_jit=False):

    def fun_tf(x):
      return 2. * x * x * x

    def fun_jax(x):
      return 3. * _maybe_jit(with_jit, jax2tf.call_tf(fun_tf))(x)

    def fun_jax_pure(x):
      return 3. * fun_tf(x)

    grad_jax = fun_jax
    grad_jax_pure = fun_jax_pure
    for _ in range(degree):
      grad_jax = jax.grad(grad_jax)
      grad_jax_pure = jax.grad(grad_jax_pure)

    res_jax = grad_jax(np.float32(5.))
    logging.info("Grad of %s degree is %s", degree, res_jax)
    self.assertAllClose(res_jax, grad_jax_pure(np.float32(5.)))

  def test_pmap(self):
    logging.info("Running test_pmap on %s devices", jax.local_device_count())

    def plus_2_tf(x):
      return tf.math.add(2., x)

    def fun_jax(x):
      return np.float32(3.) * jax2tf.call_tf(plus_2_tf)(x)

    x = np.arange(jax.local_device_count(), dtype=np.float32)
    res = jax.pmap(fun_jax)(x)
    self.assertAllClose(np.float32(3. * (x + 2)), res)

  def test_function_compile_time_constant_inputs(self):
    # Call a function for which shape inference does not give an output
    # shape.
    x = np.array([1, 2, 3], dtype=np.int32)
    def fun_tf(x):  # x:i32[3]
      # Indexing with a dynamic slice makes the TF shape inference return
      # a partially known shape.
      end_idx = x[1]
      res = x[0:end_idx]
      return res

    # Call in eager mode. Should work!
    res1 = jax2tf.call_tf(fun_tf)(x)
    self.assertAllClose(x[0:x[1]], res1)

    # Now under jit, should fail because the function is not compilable
    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
      fun_jax(x)

  def test_experimental_get_compiler_ir_design_doc(self):
    # Not a test of call_tf, but more of how experimental_get_compiler_ir works.
    # Examples are from the design doc.

    # Constant slice. This is the common case.
    x = np.zeros((10,), dtype=np.int32)

    def fun_tf(x):
      begin = 0
      return x[begin:5]

    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
    self.assertIn("(arg0.1: s32[10]) -> s32[5]", hlo)

    # Non-constant slice, but compile-time constant depending only on values.
    x = np.zeros((10,), dtype=np.int32)

    # Non-constant slice, but compile-time constant depending only on shapes.
    x = np.zeros((10,), dtype=np.int32)

    def fun_tf(x):
      begin = tf.shape(x)[0] - 2  # begin is a compile-time constant, even if x is not
      return x[begin:]

    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
    self.assertIn("(arg0.1: s32[10]) -> s32[2]", hlo)

    # Capture a variable
    outer_var = tf.Variable(np.array([3.], dtype=np.float32))
    x = np.array([2., 3., 4.], dtype=np.float32)

    def fun_tf(x):
      return x * tf.broadcast_to(outer_var, x.shape) + 1.

    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
    self.assertIn("(arg0.1: f32[3], arg1.2: f32[1]) -> f32[3]", hlo)

    # Capture a constant
    outer_ct = np.array([3.], dtype=np.float32)
    x = np.array([2., 3., 4.], dtype=np.float32)

    def fun_tf(x):
      return x * tf.broadcast_to(outer_ct, x.shape) + 1.

    hlo = tf.function(fun_tf, jit_compile=True, autograph=False).experimental_get_compiler_ir(x)()
    self.assertIn("(arg0.1: f32[3]) -> f32[3]", hlo)

    # Call get_compiler_ir in a function context
    x = np.array([2., 3., 4.], dtype=np.float32)

    def fun_tf_outer(x):
      x_const = tf.constant(0, shape=x.shape, dtype=x.dtype)
      _ = tf.function(tf.math.sin, jit_compile=True, autograph=False).experimental_get_compiler_ir(x_const)()

    # TODO(b/193754660)
    # with self.assertRaisesRegex(
    #     TypeError, "An op outside of the function building code is being passed"):
    #   tf.function(fun_tf_outer)(x)
    #
    # with self.assertRaisesRegex(
    #     TypeError, "An op outside of the function building code is being passed"):
    #   tf.function(fun_tf_outer, jit_compile=True)(x)

    # Call get_concrete_function in a graph context
    def fun_tf_outer_2(x):
      _ = tf.function(tf.math.sin, jit_compile=True).get_concrete_function(tf.TensorSpec(x.shape, x.dtype))
      return x

    # Outside of a function context, this works.
    _ = tf.function(fun_tf_outer_2)(x)
    _ = tf.function(fun_tf_outer_2, jit_compile=True)(x)

  def test_repro_193754660(self):
    # Try to reproduce b/193754660. I can't.
    # We have to have tf.function(jax2tf.convert(jax2tf.call_tf(f_tf))).
    # The get_compiler_ir will indeed fail for f_tf. Then we try to use
    # shape inference for f_tf.
    # I thought to use a f_tf that uses an op without shape inference, e.g.,
    # tfxla.gather. If we wash it through a saved_model I expect that shape
    # inference would not work on it. Instead, shape inference works!!!
    x = np.array([0, 1, 2, 3, 4, 5], dtype=np.int32)
    def f_jax(x):
      return x[1]
    f_tf = jax2tf.convert(f_jax)
    f_tf_rt, _ = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])
    f_jax2 = jax2tf.call_tf(f_tf_rt)
    f_tf2 = jax2tf.convert(f_jax2)
    res = tf.function(f_tf2, autograph=False)(x)
    self.assertAllClose(res.numpy(), f_jax(x))

  def test_effectful(self):
    x = np.ones((3,), dtype=np.float32)
    lower_effect = jax.jit(jax2tf.call_tf(tf.math.sin, has_side_effects=True)).lower(x)
    self.assertNotEmpty(lower_effect._lowering.compile_args["unordered_effects"])

    lower_no_effect = jax.jit(jax2tf.call_tf(tf.math.sin, has_side_effects=False)).lower(x)
    self.assertEmpty(lower_no_effect._lowering.compile_args["unordered_effects"])

  def test_module_documentation(self):
    def cos_tf(x):
      return tf.math.cos(x)

    # Compute cos with TF and sin with JAX
    def cos_tf_sin_jax(x):
      return jax.numpy.sin(jax2tf.call_tf(cos_tf)(x))

    # Calls `cos_tf` in TF eager mode
    x = np.float32(1.)
    cos_tf_sin_jax(x)

    # Compiles `cos_tf` using TF and embeds the XLA computation into the JAX
    # XLA computation (containing `sin`). The XLA compiler may even be able to
    # fuse through JAX-TF computations.
    jax.jit(cos_tf_sin_jax)(x)

    # Uses TF gradient for `cos_tf` and JAX gradient for `sin`
    jax.grad(cos_tf_sin_jax)(x)

    logging.info(jax.make_jaxpr(cos_tf_sin_jax)(x))
    logging.info(jax.xla_computation(cos_tf_sin_jax)(x).as_hlo_text())

  def test_tf_gather(self):
    """tf_gather gradient output is tf.IndexSlices."""
    operand = jnp.array(np.random.uniform(size=(100, 128)))
    indices = jnp.array(np.random.randint(low=0, high=100, size=(4000,)))

    @tf.function(jit_compile=True, autograph=False)
    def fun_tf(operand, indices):
      return tf.experimental.numpy.std(tf.gather(operand, indices))

    fun_jax = jax2tf.call_tf(fun_tf)
    grad_fun_jax = jax.grad(fun_jax)
    grad_res = grad_fun_jax(operand, indices)
    self.assertEqual(grad_res.shape, (100, 128))

  def test_output_shape_dtype_none(self):
    x = jnp.zeros((10), dtype=jnp.float32)

    @tf.function(jit_compile=True, autograph=False)
    def fun_tf(x):  # pylint: disable=unused-argument
      return

    fun_jax_1 = jax2tf.call_tf(fun_tf, output_shape_dtype=None)
    fun_jax_2 = jax2tf.call_tf(fun_tf)
    self.assertIsNone(fun_jax_1(x))
    self.assertIsNone(fun_jax_2(x))
    fun_jax_3 = jax2tf.call_tf(
        fun_tf, output_shape_dtype=jax.ShapeDtypeStruct((10,), jnp.float32)
    )
    with self.assertRaisesRegex(
        ValueError,
        "The pytree of the TensorFlow function results does not match the"
        " pytree of the declared output_shape_dtype",
    ):
      _ = fun_jax_3(x)

  def test_output_shape_dtype_not_none(self):
    x = jnp.zeros((10), dtype=jnp.float32)

    @tf.function(jit_compile=True, autograph=False)
    def fun_tf(x):
      return x

    fun_jax_1 = jax2tf.call_tf(
        fun_tf, output_shape_dtype=jax.ShapeDtypeStruct((10,), jnp.float32)
    )
    fun_jax_2 = jax2tf.call_tf(fun_tf)
    self.assertAllClose(fun_jax_1(x), fun_jax_2(x))

    fun_jax_3 = jax2tf.call_tf(fun_tf, output_shape_dtype=None)
    with self.assertRaisesRegex(
        ValueError,
        "The pytree of the TensorFlow function results does not match the"
        " pytree of the declared output_shape_dtype",
    ):
      _ = fun_jax_3(x)


class RoundTripToJaxTest(tf_test_util.JaxToTfTestCase):
  "Reloading output of jax2tf into JAX with call_tf"
  def setUp(self):
    if tf is None:
      raise unittest.SkipTest("Test requires tensorflow")
    # TODO(b/171320191): this line works around a missing context initialization
    # bug in TensorFlow.
    _ = tf.add(1, 1)
    super().setUp()

  def test_simple(self):
    f_jax = jnp.sin
    f_jax_rt = jax2tf.call_tf(jax2tf.convert(f_jax))
    x = np.float32(0.7)
    self.assertAllClose(f_jax(x), f_jax_rt(x))

  def test_pytree(self):
    def f_jax(x):  # x: dict(a=f32, b=f32)
      return dict(a=x["a"]+1., b=x)
    x = dict(a=0.7, b=0.8)
    f_jax_rt = jax2tf.call_tf(jax2tf.convert(f_jax))
    self.assertAllClose(f_jax(x), f_jax_rt(x))

  def test_custom_grad(self):
    @jax.custom_vjp
    def f(x):
      return x * x

    # f_fwd: a -> (b, residual)
    def f_fwd(x):
      return f(x), np.float32(3.) * x
    # f_bwd: (residual, CT b) -> [CT a]
    def f_bwd(residual, ct_b):
      return residual * ct_b,

    f.defvjp(f_fwd, f_bwd)

    f_rt = jax2tf.call_tf(jax2tf.convert(f, with_gradient=True))
    x = np.float32(0.7)
    self.assertAllClose(f(x), f_rt(x))
    self.assertAllClose(jax.grad(f)(x), jax.grad(f_rt)(x))

  def test_shape_poly(self):
    f_jax = jnp.sin
    f_jax_rt = jax2tf.call_tf(jax2tf.convert(f_jax,
                                             polymorphic_shapes=["(b, ...)"]))
    x = np.array([0.7, 0.8], dtype=np.float32)
    self.assertAllClose(f_jax(x), f_jax_rt(x))

  def test_saved_model_simple(self):
    x = np.array([0.7, 0.8], dtype=np.float32)
    def f_jax(x):
      return jnp.sin(x)

    f_tf = jax2tf.convert(f_jax)
    restored_tf, _ = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])
    restored_jax = jax2tf.call_tf(restored_tf)
    self.assertAllClose(f_jax(x), restored_jax(x))

  def test_saved_model_variables(self):
    param = np.array([1., 2.], dtype=np.float32)
    x = np.array([0.7, 0.8], dtype=np.float32)
    def f_jax(param, x):
      return jnp.sin(x) + jnp.cos(param)

    param_v = tf.Variable(param)
    f_tf = jax2tf.convert(f_jax)
    _, restored_model = tf_test_util.SaveAndLoadFunction(
        lambda x: f_tf(param_v, x),
        input_args=[x],
        variables=[param_v])
    restored_jax = jax2tf.call_tf(restored_model.f)
    self.assertAllClose(f_jax(param, x), restored_jax(x))
    self.assertAllClose(f_jax(param, x), jax.jit(restored_jax)(x))

  def test_saved_model_shape_poly(self):
    tracing_count = 0
    x = np.array([0.7, 0.8], dtype=np.float32)
    def f_jax(x):
      nonlocal tracing_count
      tracing_count += 1
      return jnp.sin(x)

    f_tf = jax2tf.convert(f_jax, polymorphic_shapes=["(b, ...)"])
    res_jax = f_jax(x)
    self.assertEqual(1, tracing_count)
    # Will trace twice, it seems. Once to get the result signature, and once again
    # for the actual saving.
    restored_f, _ = tf_test_util.SaveAndLoadFunction(
        f_tf, input_signature=[tf.TensorSpec([None], x.dtype)])
    self.assertGreaterEqual(tracing_count, 2)
    tracing_count = 0
    f_jax_rt = jax2tf.call_tf(restored_f)
    self.assertAllClose(res_jax, f_jax_rt(x))
    # Ensure that restored_f works at other batch size as well
    y = np.concatenate([x, x])
    self.assertEqual(0, tracing_count)
    res_jax_y = f_jax(y)
    self.assertEqual(1, tracing_count)
    # No more tracing for f_jax_rt
    self.assertAllClose(res_jax_y, f_jax_rt(y))
    self.assertEqual(1, tracing_count)

  def test_custom_grad_saved_model(self):

    @jax.custom_vjp
    def f(x):
      return x * x

    # f_fwd: a -> (b, residual)
    def f_fwd(x):
      return f(x), np.float32(3.) * x
    # f_bwd: (residual, CT b) -> [CT a]
    def f_bwd(residual, ct_b):
      return residual * ct_b,

    f.defvjp(f_fwd, f_bwd)
    def g(x):
      return jnp.sum(f(x))

    g_tf, _ = tf_test_util.SaveAndLoadFunction(
        jax2tf.convert(g, with_gradient=True),
        input_signature=[tf.TensorSpec(shape=(1,), dtype=tf.float32)],
    )
    g_rt = jax2tf.call_tf(g_tf)
    x = np.array([0.7], dtype=np.float32)
    self.assertAllClose(g(x), g_rt(x))
    self.assertAllClose(jax.grad(g)(x), jax.grad(g_rt)(x))

  def test_without_gradient_saved_model(self):
    # Explicitly with_gradient=False
    f_jax = jnp.sum

    x = np.array([0.7, 0.8], dtype=np.float32)
    f_tf, _ = tf_test_util.SaveAndLoadFunction(
        jax2tf.convert(f_jax, with_gradient=False),
        input_args=[x])
    f_rt = jax2tf.call_tf(f_tf)

    self.assertAllClose(f_jax(x), f_rt(x))
    with self.assertRaisesRegex(Exception,
                                "Gradient explicitly disabled.*jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients"):
      jax.grad(f_rt)(x)

  def test_saved_model_no_gradients(self):
    # Save without gradients
    f_jax = jnp.sum

    x = np.array([0.7, 0.8], dtype=np.float32)
    f_tf, _ = tf_test_util.SaveAndLoadFunction(
        jax2tf.convert(f_jax, with_gradient=True), input_args=[x],
        save_gradients=False)
    f_rt = jax2tf.call_tf(f_tf)

    self.assertAllClose(f_jax(x), f_rt(x))
    # TODO: clean this up b/191117111: it should fail with a clear error
    # The following results in a confusing error:
    # TypeError: tf.Graph captured an external symbolic tensor.
    with self.assertRaises(TypeError):
      _ = jax.grad(f_rt)(x)

  def test_call_tf_under_function_context(self):
    def fun_jax(x, y):
      z = jax2tf.call_tf(tf.math.sin)(x) + jnp.cos(y)
      return z

    x = np.array([-1.0, 0.0, 1.0], dtype=np.float32)
    y = np.array([-0.5, 0.0, 0.5], dtype=np.float32)

    converted_fun = tf.function(
        jax2tf.convert(fun_jax, native_serialization=True)
    )
    expected = np.sin(x) + np.cos(y)
    res = tf.function(converted_fun, jit_compile=True, autograph=False)(x, y)
    self.assertAllClose(expected, res.numpy(), atol=1e-5, rtol=1e-5)

  @parameterized.named_parameters(
      dict(
          testcase_name=f"_{dtype.__name__}",
          dtype=dtype,
      )
      for dtype in set(jtu.dtypes.all_floating)
  )
  def test_all_floating_input_gradient(self, dtype):
    def tf_f(x):
      res = tf.math.sin(x)
      return tf.reduce_sum(res)

    jax_f = jax2tf.call_tf(tf_f)
    tf_f_rt = jax2tf.convert(jax_f)
    x = jnp.array([5.0, 6.0, 7.0]).astype(dtype)

    def assert_all_close_support_bfloat16(baseline, candidate):
      def conversion(x):
        # convert scalar to array and bfloat16 to float32
        # to support self.assertAllClose numpy array comparison.
        if x.shape == tf.TensorShape([]):
          x = tf.convert_to_tensor([x])
        if dtype == jnp.float16:
          x = tf.cast(x, tf.float32)
        return x

      baseline = jax.tree_util.tree_map(conversion, baseline)
      candidate = jax.tree_util.tree_map(conversion, candidate)
      self.assertAllClose(baseline, candidate)

    # Eager mode
    assert_all_close_support_bfloat16(tf_f(x), tf_f_rt(x))

    # Compiled function mode
    assert_all_close_support_bfloat16(
        tf.function(tf_f)(x), tf.function(tf_f_rt)(x)
    )

    # Compiled function mode with jit_compiled=True
    assert_all_close_support_bfloat16(
        tf.function(tf_f, jit_compile=True)(x),
        tf.function(tf_f_rt, jit_compile=True)(x),
    )

    # RoundTrip test for the gradient
    grad_fun_jax = jax.grad(jax2tf.call_tf(tf_f))
    grad_fun_jax_rt = jax2tf.call_tf(jax2tf.convert(grad_fun_jax))

    # Eager mode
    assert_all_close_support_bfloat16(grad_fun_jax(x), grad_fun_jax_rt(x))

    # Jit mode
    assert_all_close_support_bfloat16(
        jax.jit(grad_fun_jax)(x), jax.jit(grad_fun_jax_rt)(x)
    )

  @parameterized.named_parameters(
      dict(
          testcase_name=f"_{dtype.__name__}",
          dtype=dtype,
      )
      for dtype in set(jtu.dtypes.complex)
  )
  def test_complex_input_gradient(self, dtype):
    def tf_f(x):
      res = tf.math.sin(x)
      return tf.reduce_sum(res)

    x = jnp.array([(5.0 + 4.0j), (6.0 + 3.0j), (7.0 + 8.0j)]).astype(dtype)

    jax_f = jax2tf.call_tf(tf_f)
    tf_f_rt = jax2tf.convert(jax_f)

    # Eager mode
    self.assertAllClose(tf_f(x), tf_f_rt(x))

    # tf.function context
    self.assertAllClose(tf.function(tf_f)(x), tf.function(tf_f_rt)(x))

    # tf.function context with jit_compiled=True
    self.assertAllClose(
        tf.function(tf_f, jit_compile=True)(x),
        tf.function(tf_f_rt, jit_compile=True)(x),
    )

    # RoundTrip test for the gradient
    grad_fun_jax = jax.grad(jax2tf.call_tf(tf_f), holomorphic=True)
    grad_fun_jax_rt = jax2tf.call_tf(jax2tf.convert(grad_fun_jax))

    # Eager mode
    self.assertAllClose(grad_fun_jax(x), grad_fun_jax_rt(x))

    # Jit mode
    self.assertAllClose(jax.jit(grad_fun_jax)(x), jax.jit(grad_fun_jax_rt)(x))


class RoundTripToTfTest(tf_test_util.JaxToTfTestCase):
  "Reloading output of call_tf into TF with jax2tf."

  def setUp(self):
    if tf is None:
      raise unittest.SkipTest("Test requires tensorflow")
    # TODO(b/171320191): this line works around a missing context initialization
    # bug in TensorFlow.
    _ = tf.add(1, 1)
    super().setUp()

  def test_alternate(self):
    # Alternate sin/cos with sin in TF and cos in JAX
    f_tf_inner = tf.math.sin
    def f_jax(x_jax):
      y_jax = jnp.cos(x_jax)
      z_jax = jax2tf.call_tf(f_tf_inner)(y_jax)
      return jnp.cos(z_jax)
    def f_tf_outer(x_tf):
      y_tf = tf.math.sin(x_tf)
      z_tf = jax2tf.convert(f_jax)(y_tf)
      return tf.math.sin(z_tf)

    x = np.float32(0.7)

    self.assertAllClose(np.sin(np.cos(np.sin(np.cos(np.sin(x))))),
                        f_tf_outer(x).numpy())
    xv = tf.Variable(x)
    with tf.GradientTape() as tape:
      res = f_tf_outer(xv)
    g_tf = tape.gradient(res, xv)
    _, gf = tf_test_util.ComputeTfValueAndGrad(f_tf_outer, (x,))
    # Eager
    expected_res = np.sin(np.cos(np.sin(np.cos(np.sin(x)))))
    self.assertAllClose(expected_res, f_tf_outer(x).numpy())

    # Gradient
    expected_grad = (np.cos(np.cos(np.sin(np.cos(np.sin(x))))) *
                     np.sin(np.sin(np.cos(np.sin(x)))) *
                     np.cos(np.cos(np.sin(x))) *
                     np.sin(np.sin(x)) *
                     np.cos(x))
    self.assertAllClose(expected_grad, g_tf.numpy())

    # Graph
    self.assertAllClose(expected_res,
                        tf.function(f_tf_outer, autograph=False)(x).numpy())

    # Compiled
    self.assertAllClose(expected_res,
                        tf.function(f_tf_outer, autograph=False,
                                    jit_compile=True)(x).numpy())

  def test_saved_model(self):
    x = np.array([.7, .8], dtype=np.float32)
    def fun_tf(x):
      return tf.math.sin(x)
    def fun_jax(x):
      return jax2tf.call_tf(fun_tf)(x)

    # Now convert and save to SavedModel
    fun_tf_rt = jax2tf.convert(fun_jax)
    res = fun_tf_rt(x)
    self.assertAllClose(np.sin(x), res.numpy())

    res = tf.function(fun_tf_rt, autograph=False)(x)
    self.assertAllClose(np.sin(x), res.numpy())

    res = tf.function(fun_tf_rt, jit_compile=True, autograph=False)(x)
    self.assertAllClose(np.sin(x), res.numpy())

    reloaded_f, _ = tf_test_util.SaveAndLoadFunction(
        fun_tf_rt, input_args=[x])
    res = reloaded_f(x)
    self.assertAllClose(np.sin(x), res.numpy())

  def test_saved_model_polymorphic_input_static_output(self):
    x = np.array([.7, .8], dtype=np.float32)
    def fun_tf(x):
      return tf.math.reduce_sum(tf.math.sin(x))
    def fun_jax(x):
      return jax2tf.call_tf(fun_tf)(x)

    # Now convert and save to SavedModel
    fun_tf_rt = jax2tf.convert(fun_jax)
    res = fun_tf_rt(x)
    self.assertAllClose(fun_tf(x), res.numpy())

    res = tf.function(fun_tf_rt, autograph=False)(x)
    self.assertAllClose(fun_tf(x), res.numpy())

    res = tf.function(fun_tf_rt, jit_compile=True, autograph=False)(x)
    self.assertAllClose(fun_tf(x), res.numpy())

    reloaded_f, _ = tf_test_util.SaveAndLoadFunction(
        fun_tf_rt, input_args=[x])
    res = reloaded_f(x)
    self.assertAllClose(fun_tf(x), res.numpy())

  def test_function_dynamic_shape(self):
    # Call a function for which shape inference does not give an output
    # shape.
    x = np.array([-1, 0, 1], dtype=np.int32)
    def fun_tf(x):  # x:i32[3]
      # The shape depends on the value of x
      return tf.cond(x[0] >= 0, lambda: x, lambda: x[1:])

    # Call in eager mode. Should work!
    res1 = jax2tf.call_tf(fun_tf)(x)
    expected = x[1:]
    self.assertAllClose(expected, res1, check_dtypes=False)

    # Now under jit, should fail because the function is not compilable
    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
      fun_jax = jax.jit(jax2tf.call_tf(fun_tf))
      fun_jax(x)

    # TODO(necula): this should work in op-by-op mode, but it fails because
    # jax2tf.convert does abstract evaluation.
    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
      fun_tf_rt = jax2tf.convert(jax2tf.call_tf(fun_tf))
      fun_tf_rt(x)

  @_parameterized_jit
  def test_shape_poly_static_output_shape(self, with_jit=True):
    if config.jax2tf_default_native_serialization:
      raise unittest.SkipTest("TODO(b/268386622): call_tf with shape polymorphism and native serialization.")
    x = np.array([0.7, 0.8], dtype=np.float32)

    def fun_tf(x):
      return tf.math.reduce_sum(tf.math.sin(x))

    fun_jax = jax2tf.call_tf(fun_tf)
    fun_tf_rt = _maybe_tf_jit(with_jit,
        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
    self.assertAllClose(fun_tf(x), fun_tf_rt(x))

  @_parameterized_jit
  def test_shape_poly(self, with_jit=False):
    if config.jax2tf_default_native_serialization:
      raise unittest.SkipTest("TODO(b/268386622): call_tf with shape polymorphism and native serialization.")
    x = np.array([7, 8, 9, 10], dtype=np.float32)
    def fun_jax(x):
      y = jax2tf.call_tf(tf.math.sin,
                         output_shape_dtype=jax.ShapeDtypeStruct(x.shape, x.dtype))(x)
      z = jnp.cos(y)
      w = jax2tf.call_tf(lambda z: tf.concat([z, z], axis=0),
                         output_shape_dtype=jax.ShapeDtypeStruct((2 * z.shape[0],), z.dtype))(z)
      assert w.shape[0] == 2 * x.shape[0]
      return w

    fun_tf_rt = _maybe_tf_jit(with_jit,
        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
    res_tf = fun_tf_rt(x)
    self.assertAllClose(fun_jax(x), res_tf)

  @_parameterized_jit
  def test_shape_poly_pytree_result(self, with_jit=True):
    if config.jax2tf_default_native_serialization:
      raise unittest.SkipTest("TODO(b/268386622): call_tf with shape polymorphism and native serialization.")
    x = np.array([7, 8, 9, 10], dtype=np.float32)
    def fun_jax(x):
      # Returns a tuple
      y = jax2tf.call_tf(lambda x: (x, tf.concat([x, x], axis=0)),
          output_shape_dtype=(jax.ShapeDtypeStruct(x.shape, x.dtype),
                              jax.ShapeDtypeStruct((2 * x.shape[0],), x.dtype)))(x)
      assert y[0].shape[0] == x.shape[0]
      assert y[1].shape[0] == 2 * x.shape[0]
      return y

    fun_tf_rt = _maybe_tf_jit(with_jit,
        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
    res_tf = fun_tf_rt(x)
    self.assertAllClose(fun_jax(x), res_tf)

  @_parameterized_jit
  def test_shape_poly_error_no_output_shape_dtype(self, with_jit=True):
    x = np.array([7, 8, 9, 10], dtype=np.float32)
    def fun_jax(x):
      return jax2tf.call_tf(tf.math.sin)(x)

    fun_tf_rt = _maybe_tf_jit(with_jit,
        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
    with self.assertRaisesRegex(ValueError, _call_tf_dynamic_shape_error):
      fun_tf_rt(x)

  @_parameterized_jit
  def test_shape_poly_error_mismatch_output_shape_dtype_tree(self, with_jit=False):
    x = np.array([7, 8, 9, 10], dtype=np.float32)
    def fun_jax(x):
      return jax2tf.call_tf(tf.math.sin,
          output_shape_dtype=(jax.ShapeDtypeStruct(x.shape, x.dtype),
                              jax.ShapeDtypeStruct(x.shape, x.dtype)))(x)

    fun_tf_rt = _maybe_tf_jit(with_jit,
        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))

    with self.assertRaisesRegex(
        ValueError,
        "The pytree of the TensorFlow function results does not match the pytree of the declared output_shape_dtype"):
      fun_tf_rt(x)

  @parameterized.named_parameters(
      _named_test(with_jit=with_jit, kind=kind)
      for with_jit in [True, False]
      for kind in ["bad_rank", "bad_dim", "bad_dtype", "bad_dtype_x64"])
  def test_shape_poly_error_mismatch_output_shape_dtype(self, with_jit=False, kind="bad_rank"):
    x = np.array([7, 8, 9, 10], dtype=np.float32)

    if kind == "bad_rank":
      def fun_jax(x):
        return jax2tf.call_tf(lambda x: x,
                              # Wrong shape rank
                              output_shape_dtype=jax.ShapeDtypeStruct((), x.dtype))(x)
    elif kind == "bad_dim":
      def fun_jax(x):
        bad_shape = (5 + x.shape[0],)
        y = jax2tf.call_tf(lambda x: x,
                           # Wrong dimension
                           output_shape_dtype=jax.ShapeDtypeStruct(bad_shape, x.dtype))(x)
        # JAX will believe that the following is Ok, leading to downstream error in TF
        return y + jnp.ones(bad_shape, dtype=x.dtype)
    elif kind == "bad_dtype":
      def fun_jax(x):
        return jax2tf.call_tf(lambda x: x,
                              output_shape_dtype=jax.ShapeDtypeStruct(x.shape, np.int32))(x)
    elif kind == "bad_dtype_x64":
      def fun_jax(x):
        return jax2tf.call_tf(lambda x: x * np.float64(3.),
                              output_shape_dtype=jax.ShapeDtypeStruct(x.shape, np.float64))(x)
    else:
      assert False
    expect_ex = ValueError
    expect_error = r"The shapes or dtypes returned by the TensorFlow function do not match the declared output_shape_dtype"

    # Call without shape polymorphism
    fun_tf_rt = _maybe_tf_jit(with_jit, jax2tf.convert(fun_jax))
    with self.assertRaisesRegex(expect_ex, expect_error):
      fun_tf_rt(x)

    # Now with shape polymorphism
    if kind == "bad_dim" and with_jit:
      # TODO: in jit more the error pops up later, at AddV2
      expect_error = "Dimensions must be equal, but are 4 and 9 for .* AddV2"
    if kind == "bad_dim" and config.jax2tf_default_native_serialization:
      # TODO(b/268386622): call_tf with shape polymorphism and native serialization.
      expect_error = "Error compiling TensorFlow function"
    fun_tf_rt = _maybe_tf_jit(with_jit,
        jax2tf.convert(fun_jax, polymorphic_shapes=["b, ..."]))
    with self.assertRaisesRegex(expect_ex, expect_error):
      fun_tf_rt(x)

  def test_inner_native_serialization(self):
    # Two nested jax2tf, the inner one being with native serialization
    x = np.ones((3,), dtype=np.float32)
    def f_inner_jax(x):
      return jnp.sin(x)
    def f_outer_jax(x):
      f_inner_tf = jax2tf.convert(f_inner_jax, native_serialization=True)
      return jnp.cos(jax2tf.call_tf(f_inner_tf)(x))

    f_outer_tf = tf.function(
        jax2tf.convert(f_outer_jax, native_serialization=False),
        autograph=False)
    f_outer_graph = str(f_outer_tf.get_concrete_function(tf.convert_to_tensor(x)).graph.as_graph_def())
    # Quick way to check that there is an XlaCallModule op, and a Cos op, but no Sin op
    self.assertIn('op: "Cos"', f_outer_graph)
    self.assertIn('op: "XlaCallModule"', f_outer_graph)
    self.assertNotIn('op: "Sin"', f_outer_graph)

  @parameterized.named_parameters(
      _named_test(f2_function=f2_function, f2_saved_model=f2_saved_model,
                  f4_function=f4_function, f4_saved_model=f4_saved_model)
      for f2_function in [True, False]
      for f2_saved_model in [True, False]
      for f4_function in [True, False]
      for f4_saved_model in [True, False])
  def test_several_round_trips(self,
                               f2_function=False, f2_saved_model=False,
                               f4_function=False, f4_saved_model=False):
    x = np.array(.7, dtype=np.float32)
    # f(n)(x) = 2. * x^n
    def f(n):
      def fn(x):
        acc = np.array(2., dtype=x.dtype)
        for i in range(n):
          acc *= x
        return acc
      return fn

    f2_tf = lambda x: x * jax2tf.convert(f(1))(x)
    if f2_function:
      f2_tf = tf.function(f2_tf, autograph=False)
    if f2_saved_model:
      f2_tf, _ = tf_test_util.SaveAndLoadFunction(f2_tf, input_args=[x])

    self.assertAllClose(f(2)(x), f2_tf(x).numpy())
    _, (g_f2_ft,) = tf_test_util.ComputeTfValueAndGrad(f2_tf, [x])
    self.assertAllClose(jax.grad(f(2))(x), g_f2_ft.numpy())

    f3_jax = lambda x: x * jax2tf.call_tf(f2_tf)(x)
    self.assertAllClose(f(3)(x), f3_jax(x))
    self.assertAllClose(f(3)(x), jax.jit(f3_jax)(x))
    self.assertAllClose(jax.grad(f(3))(x), jax.grad(f3_jax)(x))

    f4_tf = lambda x: x * jax2tf.convert(f3_jax)(x)
    self.assertAllClose(f(4)(x), f4_tf(x).numpy())
    _, (g_f4_ft,) = tf_test_util.ComputeTfValueAndGrad(f4_tf, [x])
    self.assertAllClose(jax.grad(f(4))(x), g_f4_ft.numpy())

    if f4_function:
      f4_tf = tf.function(f4_tf, autograph=False)
    if f4_saved_model:
      f4_tf, _ = tf_test_util.SaveAndLoadFunction(f4_tf, input_args=[x])
    self.assertAllClose(f(4)(x), f4_tf(x).numpy())
    _, (g_f4_ft,) = tf_test_util.ComputeTfValueAndGrad(f4_tf, [x])
    self.assertAllClose(jax.grad(f(4))(x), g_f4_ft.numpy())

  @classmethod
  def _walk_stablehlo_operations(cls, op, cb):
    """walk the stablehlo operation recursive with callback function."""
    cb(op)
    for region in op.operation.regions:
      for block in region:
        for op in block:
          cls._walk_stablehlo_operations(op, cb)

  def test_call_tf_graph(self):
    const = tf.Variable(0.0, dtype=tf.float32)

    @tf.function(jit_compile=True)
    def tf_func_1(x):
      return x * x + const

    @tf.function
    def tf_func_2(x, y):
      return tf_func_1(x) + y

    @tf.function
    def tf_func_3(x, y, z):
      return tf_func_2(x, y) + z, z

    x = jnp.array(3.0, dtype=jnp.float32)
    y = jnp.array(3.0, dtype=jnp.float32)
    z = jnp.array(5.0, dtype=jnp.float32)
    f_jax = jax.jit(jax2tf.call_tf(tf_func_3, call_tf_graph=False))
    stablehlo_module = f_jax.lower(x, y, z).compiler_ir("stablehlo")
    self.assertNotIn("stablehlo.custom_call", str(stablehlo_module))

    f_jax = jax.jit(
        jax2tf.call_tf(
            tf_func_3,
            call_tf_graph=True,
        )
    )
    with self.assertRaisesRegex(
        ValueError,
        "call_tf_graph=True only support exporting by jax2tf.convert currently",
    ):
      stablehlo_module = f_jax.lower(x, y, z).compiler_ir("stablehlo")
      self.assertIn("stablehlo.custom_call", str(stablehlo_module))

      called_index_list = []

      def _extract_info(op):
        if op.operation.name != "stablehlo.custom_call":
          return
        tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
        called_index = ir.IntegerAttr(tf_backend_config["called_index"]).value
        called_index_list.append(called_index)

      self._walk_stablehlo_operations(stablehlo_module, _extract_info)
      self.assertLen(called_index_list, 1)

  @parameterized.named_parameters(
      dict(
          testcase_name="multiple_outputs",
          tf_f=lambda x: tf.py_function(np.sin, [x], tf.float32),
          output_shape_dtype=jax.ShapeDtypeStruct((10,), jnp.float32),
      ),
      dict(
          testcase_name="zero_outputs",
          tf_f=lambda x: print(tf.strings.length(tf.constant("hello, world"))),
          output_shape_dtype=None,
      ),
  )
  def test_call_tf_graph_non_compilable(self, tf_f, output_shape_dtype):
    inputs = jnp.ones([10], dtype=jnp.float32)
    called_index_list = []
    xla_call_module_list = []

    def _extract_info(op):
      if op.operation.name != "stablehlo.custom_call":
        return
      tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
      called_index = ir.IntegerAttr(tf_backend_config["called_index"]).value
      called_index_list.append(called_index)

    jax_f = jax2tf.call_tf(
        tf_f,
        call_tf_graph=True,
        output_shape_dtype=output_shape_dtype,
    )

    # Eager mode
    self.assertAllClose(tf_f(inputs), jax_f(inputs))

    # Jit mode
    stablehlo_module = None
    with self.assertRaisesRegex(
        ValueError,
        "call_tf_graph=True only support exporting by jax2tf.convert currently",
    ):
      stablehlo_module = jax.jit(jax_f).lower(inputs).compiler_ir("stablehlo")
    if stablehlo_module:
      self.assertIn(
          "stablehlo.custom_call @tf.call_tf_function",
          str(stablehlo_module),
      )
      self.assertIn("tf.backend_config", str(stablehlo_module))
      self._walk_stablehlo_operations(stablehlo_module, _extract_info)
      self.assertLen(called_index_list, 1)

    # Test model exporting and reloading.
    # There is no runtime support yet so it can not run.
    tf_f_rt = jax2tf.convert(
        jax_f,
        native_serialization=True,
        with_gradient=False,
    )
    _, restored_model = tf_test_util.SaveAndLoadFunction(
        tf_f_rt, input_args=[inputs]
    )
    func_def = restored_model.f.concrete_functions[0]

    for node_def in func_def.graph.as_graph_def().node:
      if node_def.op == "XlaCallModule":
        xla_call_module_list.append(node_def)
    # There is only one xla_call_module in the saved model.
    self.assertLen(xla_call_module_list, 1)

    # Check the xla_call_module version and function_list attributes.
    xla_call_module = xla_call_module_list[0]
    self.assertGreaterEqual(xla_call_module.attr["version"].i, 5)
    self.assertIn("function_list", str(xla_call_module.attr))
    xla_call_module_list.clear()
    called_index_list.clear()

    # If JAX calls same tensorflow function by `jax2tf.call_tf` twice,
    # it should return two different tf concrete functions.
    def jax_f_2(x):
      res1 = jax2tf.call_tf(
          tf_f,
          call_tf_graph=True,
          output_shape_dtype=output_shape_dtype,
      )(x)
      res2 = jax2tf.call_tf(
          tf_f,
          call_tf_graph=True,
          output_shape_dtype=output_shape_dtype,
      )(x)
      return res1, res2
    stablehlo_module = None
    with self.assertRaisesRegex(ValueError, "call_tf_graph=True only support exporting by jax2tf.convert currently"):
      stablehlo_module = jax.jit(jax_f_2).lower(inputs).compiler_ir("stablehlo")
    if stablehlo_module:
      self._walk_stablehlo_operations(stablehlo_module, _extract_info)
    xla_call_module_list.clear()

  def test_b279454591(self):
    """Test case when tensorflow function returns `StatefulPartitionedCall` op."""
    inputs = jnp.ones([10], dtype=jnp.float32)

    # With one or more outputs, it is okay.
    def tf_f(x):
      y = tf.math.sin(3.0)
      tf.print(y)
      return x

    jax_f = jax2tf.call_tf(
        tf.function(tf_f),
        call_tf_graph=True,
    )
    tf_f_rt = jax2tf.convert(
        jax_f,
        native_serialization=True,
        with_gradient=False,
    )
    _, _ = tf_test_util.SaveAndLoadFunction(tf_f_rt, input_args=[inputs])

    # With zero output, it return `StatefulPartitionedCall` op instead.
    def tf_f_2():
      y = tf.math.sin(3.0)
      tf.print(y)
      return

    jax_f_2 = jax2tf.call_tf(tf.function(tf_f_2), call_tf_graph=True)
    tf_f_rt_2 = jax2tf.convert(
        jax_f_2,
        native_serialization=True,
        with_gradient=False,
    )
    _, _ = tf_test_util.SaveAndLoadFunction(tf_f_rt_2, input_args=[])

  def test_call_tf_graph_ordered(self):
    @tf.function
    def tf_print(x):
      tf.print(x)

    call_tf_print = jax2tf.call_tf(
        tf_print,
        call_tf_graph=True,
        ordered=True,
    )

    x = jnp.array(1.0, dtype=jnp.float32)

    def body(i, x):
      call_tf_print(x)
      return x + 1

    @jax.jit
    def f_jax(x):
      return jax.lax.fori_loop(0, 4, body, x)

    num_custom_calls = 0

    def _check_mlir_ops(op):
      nonlocal num_custom_calls

      if (
          op.operation.name == "stablehlo.custom_call"
          and ir.StringAttr(op.attributes["call_target_name"]).value
          == "tf.call_tf_function"
      ):
        num_custom_calls += 1

        # The custom call op must have `has_token_input_output` attribute.
        tf_backend_config = ir.DictAttr(op.attributes["tf.backend_config"])
        self.assertTrue(
            ir.BoolAttr(tf_backend_config["has_token_input_output"]).value
        )

        # Verify that the first argument/result of the custom call op is a token
        # type. This is a calling convention defined by `has_token_input_output`.
        self.assertTrue(hlo.TokenType.isinstance(op.operands[0].type))
        self.assertTrue(hlo.TokenType.isinstance(op.results[0].type))

    stablehlo_module = None
    with self.assertRaisesRegex(
        ValueError,
        "call_tf_graph=True only support exporting by jax2tf.convert currently",
    ):
      lower = f_jax.lower(x)
      self.assertNotEmpty(lower._lowering.compile_args["ordered_effects"])
      stablehlo_module = lower.compiler_ir("stablehlo")
    if stablehlo_module:
      self._walk_stablehlo_operations(stablehlo_module, _check_mlir_ops)
      self.assertEqual(num_custom_calls, 1)

    f_tf = jax2tf.convert(
        f_jax,
        native_serialization=True,
        with_gradient=False,
    )
    _, restored_model = tf_test_util.SaveAndLoadFunction(f_tf, input_args=[x])

  @parameterized.named_parameters([
    dict(testcase_name=f"{ordered=}", ordered=ordered)
    for ordered in [True, False]
  ])
  def test_call_tf_graph_polymorphic(self, ordered: bool):
    @tf.function(jit_compile=True, autograph=False)
    @partial(jax2tf.convert,
      with_gradient=False,
      native_serialization=True,
      polymorphic_shapes=["(b)"])
    @jax.jit
    def tf_f_2(x):
      tf_f = lambda x: print(tf.strings.length(tf.constant("hello, world")))
      jax2tf.call_tf(tf_f,
                     call_tf_graph=True,
                     ordered=ordered)(x)
      return x

    x = np.arange(3, dtype=np.int32)
    _ = tf.function(tf_f_2, autograph=False).get_concrete_function(x)

  # TODO(b/293927250): call_tf_graph=True only accept concrete_function. The
  # workaround here is to set `module.call=concrete_fn.`.
  @unittest.skip(
      "The root cause here is because the XLACallModule.function_list attribute"
      " depends on JAX call_tf lowering. The 2nd time tf.SavedModel TF tracing"
      " will not trigger call_tf tracing since it was already cached. The"
      " solution is to create the `CallTFContext` to make TF tracing and JAX"
      " tracing work together correctly."
  )
  def test_call_tf_graph_save_and_load(self):
    def jax_func(x):
      def func_tf(x):
        return tf.math.sin(x)

      return jnp.cos(
          jax2tf.call_tf(func_tf, output_shape_dtype=x, call_tf_graph=True)(x)
      )
    data_inputs = (np.array([0.5, 0.7], dtype=np.float32),)

    def tf_func(the_input):
      res = jax2tf.convert(jax_func, native_serialization=True)(the_input)
      return tf.identity(res, name="the_result")

    jit_tf_func = tf.function(
        tf_func,
        autograph=False,
        jit_compile=True,
    )
    # The next line is necessary to reproduce this issue. It trigger TF
    # ConcreteFunction tracing. Otherwise, you will fail with another error
    # `Found zero restored functions for caller function`.
    _ = jit_tf_func.get_concrete_function(*data_inputs)
    module = tf.Module()
    module.call = jit_tf_func # Switching to concrete_function works.
    root_dir = self.create_tempdir()
    saved_model_dir = os.path.join(root_dir, "saved_model")
    tf.saved_model.save(
        module,
        saved_model_dir,
        options=tf.saved_model.SaveOptions(experimental_custom_gradients=False),
    )
    loaded_model = tf.saved_model.load(saved_model_dir)
    res = loaded_model.call(*data_inputs)
    self.assertAllClose(jax_func(*data_inputs), res)

if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())
