@functools.partial(api_util.api_hook, tag='jax2tf_convert')
def convert(fun: Callable, *, polymorphic_shapes: Optional[Sequence[Any]]=None, with_gradient=True, enable_xla=True) -> Callable:
    'Transforms `fun` to be executed by TensorFlow.\n\n  See\n  [README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md)\n  for more details about usage and common problems.\n\n  Args:\n    fun: Function to be transformed. Its arguments and return value should be\n      JAX arrays, or nested standard Python containers (tuple/list/dict) thereof\n      (pytrees).\n    polymorphic_shapes: Specifies input shapes to be treated polymorphically\n      during conversion.\n      .. warning:: The shape-polymorphic conversion is an experimental feature.\n        It is meant to be sound, but it is known to reject some JAX programs\n        that are shape polymorphic. The details of this feature can change.  It\n        should be a Python object with the same pytree structure as, or a prefix\n        of, the tuple of arguments to the function, but with a shape\n        specification corresponding to each argument. The default value is\n        `None`, which is a shortcut for a tuple of `None` one for each argument,\n        denoting that all shapes are monomorphic.\n      See [how optional parameters are matched to\n        arguments](https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees).\n        A shape specification for an array argument should be an object\n        `PolyShape(dim0, dim1, ..., dimn)`\n      where each `dim` is a dimension specification: a positive integer denoting\n        a monomorphic dimension of the given size, or a string denoting a\n        dimension variable assumed to range over non-zero dimension sizes, or\n        the special placeholder string "_" denoting a monomorphic dimension\n        whose size is given by the actual argument. As a shortcut, an Ellipsis\n        suffix in the list of dimension specifications stands for a list of "_"\n        placeholders. For convenience, a shape specification can also be given\n        as a string\n      representation, e.g.: "batch, ...", "batch, height, width, _", possibly\n      with surrounding parentheses: "(batch, ...)".\n\n      The conversion fails if it cannot ensure that the it would produce the same\n      sequence of TF ops for any non-zero values of the dimension variables.\n\n      polymorphic_shapes are only supported for positional arguments; shape\n      polymorphism is not supported for keyword arguments.\n\n      See [the README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md#shape-polymorphic-conversion)\n      for more details.\n\n    in_shapes: DEPRECATED in favor of `polymorphic_shapes`.\n    with_gradient: if set, will add a tf.custom_gradient to the converted\n      function, by converting the ``jax.vjp(fun)``. Only first-order\n      differentiation is supported for now. If the converted function is saved\n      in a SavedModel, the custom gradients are currently lost and an error will\n      be raised if a gradient computation is attempted. This is due to a current\n      bug in TensorFlow.\n    enable_xla: if unset, the converter will try harder to use pure TF ops to\n      convert the function, and raise an error if it can not be converted\n      without resorting to XLA ops (default: True).\n\n  Returns:\n    A version of `fun` that expects TfVals as arguments (or\n    tuple/lists/dicts) thereof, and returns TfVals as outputs.\n  '
    api._check_callable(fun)

    def converted_fun(*args: TfVal, **kwargs: TfVal) -> TfVal:
        if (not core.trace_state_clean()):
            raise ValueError(('convert must be used outside all JAX transformations.' + f'Trace state: {core.thread_local_state.trace_state}'))

        def check_arg(a):
            if (not _is_tfval(a)):
                msg = f'Argument {a} of type {type(a)} of jax2tf.convert(f) should be NumPy array, scalar, tf.Variable, or tf.Tensor'
                raise TypeError(msg)
        tree_util.tree_map(check_arg, args)
        tree_util.tree_map(check_arg, list(kwargs.values()))
        args = tuple((tree_util.tree_map((lambda x, i=i: tf.identity(x, f'jax2tf_arg_{i}')), a) for (i, a) in enumerate(args)))
        kwargs = {k: tf.identity(v, f'jax2tf_arg_{k}') for (k, v) in kwargs.items()}
        (args_flat, in_tree) = tree_util.tree_flatten((args, kwargs))
        if (polymorphic_shapes is None):
            polymorphic_shapes_ = ((None,) * len(args))
        else:
            if ((not isinstance(polymorphic_shapes, Sequence)) or (len(args) != len(polymorphic_shapes))):
                msg = f'polymorphic_shapes must be a sequence with the same length as the positional argument list ({len(args)}). Got polymorphic_shapes={polymorphic_shapes}.'
                raise TypeError(msg)
            polymorphic_shapes_ = tuple(polymorphic_shapes)
        polymorphic_shapes_flat = tuple(api_util.flatten_axes('jax2tf.convert polymorphic_shapes', in_tree.children()[0], polymorphic_shapes_))
        polymorphic_shapes_flat = (polymorphic_shapes_flat + tuple(((None,) * (len(args_flat) - len(polymorphic_shapes_flat)))))
        (args_avals_flat, shapeenv) = _args_to_avals_and_env(args_flat, polymorphic_shapes_flat)
        f = lu.wrap_init(fun)
        (flat_fun, out_tree_thunk) = flatten_fun(f, in_tree)

        def converted_grad_fn(*out_cts_flat: TfVal, _out_cts_avals: Sequence[core.AbstractValue], variables=None):
            if variables:
                raise ValueError(f'Unexpected variables used in forward pass. This should not happen for first-order differentiation. variables={variables}')

            def fun_vjp_jax(args_jax, out_cts_jax):
                (_, pullback_jax) = jax.vjp(fun, *args_jax)
                return pullback_jax(out_cts_jax)
            if (polymorphic_shapes is None):
                vjp_polymorphic_shapes = None
            else:
                args_polymorphic_shapes = tree_util.tree_unflatten(in_tree.children()[0], polymorphic_shapes_flat)
                out_cts_polymorphic_shapes = tree_util.tree_unflatten(out_tree_thunk(), tuple((str(out_aval.shape) for out_aval in _out_cts_avals)))
                vjp_polymorphic_shapes = [args_polymorphic_shapes, out_cts_polymorphic_shapes]
            out_cts = tree_util.tree_unflatten(out_tree_thunk(), out_cts_flat)
            with tf.name_scope('jax2tf_vjp'):
                in_cts = convert(fun_vjp_jax, with_gradient=False, polymorphic_shapes=vjp_polymorphic_shapes)(args, out_cts)
            return in_cts
        try:
            global _shape_env
            assert (not _shape_env), f'Unexpected shape environment {_shape_env}'
            global _enable_xla
            prev_enable_xla = _enable_xla
            _enable_xla = enable_xla
            _shape_env = shapeenv
            if with_gradient:

                @tf.custom_gradient
                def converted_fun_flat_with_custom_gradient(*args_flat: TfVal) -> TfVal:
                    out_with_avals = _interpret_fun(flat_fun, args_flat, args_avals_flat)
                    (outs, out_avals) = util.unzip2(out_with_avals)
                    return (tuple(outs), functools.partial(converted_grad_fn, _out_cts_avals=tuple(out_avals)))
                out_flat = converted_fun_flat_with_custom_gradient(*args_flat)
            else:
                out_flat_raw = _interpret_fun(flat_fun, args_flat, args_avals_flat)
                message = 'The jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients'
                out_flat = [tf.raw_ops.PreventGradient(input=o, message=message) for (o, _) in out_flat_raw]
        finally:
            _shape_env = {}
            _enable_xla = prev_enable_xla
        out_flat = [tf.identity(x, 'jax2tf_out') for x in out_flat]
        out = tree_util.tree_unflatten(out_tree_thunk(), out_flat)
        return out
    return converted_fun
Transforms `fun` to be executed by TensorFlow.

See
[README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md)
for more details about usage and common problems.

Args:
  fun: Function to be transformed. Its arguments and return value should be
    JAX arrays, or nested standard Python containers (tuple/list/dict) thereof
    (pytrees).
  polymorphic_shapes: Specifies input shapes to be treated polymorphically
    during conversion.
    .. warning:: The shape-polymorphic conversion is an experimental feature.
      It is meant to be sound, but it is known to reject some JAX programs
      that are shape polymorphic. The details of this feature can change.  It
      should be a Python object with the same pytree structure as, or a prefix
      of, the tuple of arguments to the function, but with a shape
      specification corresponding to each argument. The default value is
      `None`, which is a shortcut for a tuple of `None` one for each argument,
      denoting that all shapes are monomorphic.
    See [how optional parameters are matched to
      arguments](https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees).
      A shape specification for an array argument should be an object
      `PolyShape(dim0, dim1, ..., dimn)`
    where each `dim` is a dimension specification: a positive integer denoting
      a monomorphic dimension of the given size, or a string denoting a
      dimension variable assumed to range over non-zero dimension sizes, or
      the special placeholder string "_" denoting a monomorphic dimension
      whose size is given by the actual argument. As a shortcut, an Ellipsis
      suffix in the list of dimension specifications stands for a list of "_"
      placeholders. For convenience, a shape specification can also be given
      as a string
    representation, e.g.: "batch, ...", "batch, height, width, _", possibly
    with surrounding parentheses: "(batch, ...)".

    The conversion fails if it cannot ensure that the it would produce the same
    sequence of TF ops for any non-zero values of the dimension variables.

    polymorphic_shapes are only supported for positional arguments; shape
    polymorphism is not supported for keyword arguments.

    See [the README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md#shape-polymorphic-conversion)
    for more details.

  in_shapes: DEPRECATED in favor of `polymorphic_shapes`.
  with_gradient: if set, will add a tf.custom_gradient to the converted
    function, by converting the ``jax.vjp(fun)``. Only first-order
    differentiation is supported for now. If the converted function is saved
    in a SavedModel, the custom gradients are currently lost and an error will
    be raised if a gradient computation is attempted. This is due to a current
    bug in TensorFlow.
  enable_xla: if unset, the converter will try harder to use pure TF ops to
    convert the function, and raise an error if it can not be converted
    without resorting to XLA ops (default: True).

Returns:
  A version of `fun` that expects TfVals as arguments (or
  tuple/lists/dicts) thereof, and returns TfVals as outputs.
@functools.partial(api_util.api_hook, tag='jax2tf_convert')
def convert(fun: Callable, *, polymorphic_shapes: Optional[Sequence[Any]]=None, with_gradient=True, enable_xla=True) -> Callable:
    'Transforms `fun` to be executed by TensorFlow.\n\n  See\n  [README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md)\n  for more details about usage and common problems.\n\n  Args:\n    fun: Function to be transformed. Its arguments and return value should be\n      JAX arrays, or nested standard Python containers (tuple/list/dict) thereof\n      (pytrees).\n    polymorphic_shapes: Specifies input shapes to be treated polymorphically\n      during conversion.\n      .. warning:: The shape-polymorphic conversion is an experimental feature.\n        It is meant to be sound, but it is known to reject some JAX programs\n        that are shape polymorphic. The details of this feature can change.  It\n        should be a Python object with the same pytree structure as, or a prefix\n        of, the tuple of arguments to the function, but with a shape\n        specification corresponding to each argument. The default value is\n        `None`, which is a shortcut for a tuple of `None` one for each argument,\n        denoting that all shapes are monomorphic.\n      See [how optional parameters are matched to\n        arguments](https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees).\n        A shape specification for an array argument should be an object\n        `PolyShape(dim0, dim1, ..., dimn)`\n      where each `dim` is a dimension specification: a positive integer denoting\n        a monomorphic dimension of the given size, or a string denoting a\n        dimension variable assumed to range over non-zero dimension sizes, or\n        the special placeholder string "_" denoting a monomorphic dimension\n        whose size is given by the actual argument. As a shortcut, an Ellipsis\n        suffix in the list of dimension specifications stands for a list of "_"\n        placeholders. For convenience, a shape specification can also be given\n        as a string\n      representation, e.g.: "batch, ...", "batch, height, width, _", possibly\n      with surrounding parentheses: "(batch, ...)".\n\n      The conversion fails if it cannot ensure that the it would produce the same\n      sequence of TF ops for any non-zero values of the dimension variables.\n\n      polymorphic_shapes are only supported for positional arguments; shape\n      polymorphism is not supported for keyword arguments.\n\n      See [the README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md#shape-polymorphic-conversion)\n      for more details.\n\n    in_shapes: DEPRECATED in favor of `polymorphic_shapes`.\n    with_gradient: if set, will add a tf.custom_gradient to the converted\n      function, by converting the ``jax.vjp(fun)``. Only first-order\n      differentiation is supported for now. If the converted function is saved\n      in a SavedModel, the custom gradients are currently lost and an error will\n      be raised if a gradient computation is attempted. This is due to a current\n      bug in TensorFlow.\n    enable_xla: if unset, the converter will try harder to use pure TF ops to\n      convert the function, and raise an error if it can not be converted\n      without resorting to XLA ops (default: True).\n\n  Returns:\n    A version of `fun` that expects TfVals as arguments (or\n    tuple/lists/dicts) thereof, and returns TfVals as outputs.\n  '
    api._check_callable(fun)

    def converted_fun(*args: TfVal, **kwargs: TfVal) -> TfVal:
        if (not core.trace_state_clean()):
            raise ValueError(('convert must be used outside all JAX transformations.' + f'Trace state: {core.thread_local_state.trace_state}'))

        def check_arg(a):
            if (not _is_tfval(a)):
                msg = f'Argument {a} of type {type(a)} of jax2tf.convert(f) should be NumPy array, scalar, tf.Variable, or tf.Tensor'
                raise TypeError(msg)
        tree_util.tree_map(check_arg, args)
        tree_util.tree_map(check_arg, list(kwargs.values()))
        args = tuple((tree_util.tree_map((lambda x, i=i: tf.identity(x, f'jax2tf_arg_{i}')), a) for (i, a) in enumerate(args)))
        kwargs = {k: tf.identity(v, f'jax2tf_arg_{k}') for (k, v) in kwargs.items()}
        (args_flat, in_tree) = tree_util.tree_flatten((args, kwargs))
        if (polymorphic_shapes is None):
            polymorphic_shapes_ = ((None,) * len(args))
        else:
            if ((not isinstance(polymorphic_shapes, Sequence)) or (len(args) != len(polymorphic_shapes))):
                msg = f'polymorphic_shapes must be a sequence with the same length as the positional argument list ({len(args)}). Got polymorphic_shapes={polymorphic_shapes}.'
                raise TypeError(msg)
            polymorphic_shapes_ = tuple(polymorphic_shapes)
        polymorphic_shapes_flat = tuple(api_util.flatten_axes('jax2tf.convert polymorphic_shapes', in_tree.children()[0], polymorphic_shapes_))
        polymorphic_shapes_flat = (polymorphic_shapes_flat + tuple(((None,) * (len(args_flat) - len(polymorphic_shapes_flat)))))
        (args_avals_flat, shapeenv) = _args_to_avals_and_env(args_flat, polymorphic_shapes_flat)
        f = lu.wrap_init(fun)
        (flat_fun, out_tree_thunk) = flatten_fun(f, in_tree)

        def converted_grad_fn(*out_cts_flat: TfVal, _out_cts_avals: Sequence[core.AbstractValue], variables=None):
            if variables:
                raise ValueError(f'Unexpected variables used in forward pass. This should not happen for first-order differentiation. variables={variables}')

            def fun_vjp_jax(args_jax, out_cts_jax):
                (_, pullback_jax) = jax.vjp(fun, *args_jax)
                return pullback_jax(out_cts_jax)
            if (polymorphic_shapes is None):
                vjp_polymorphic_shapes = None
            else:
                args_polymorphic_shapes = tree_util.tree_unflatten(in_tree.children()[0], polymorphic_shapes_flat)
                out_cts_polymorphic_shapes = tree_util.tree_unflatten(out_tree_thunk(), tuple((str(out_aval.shape) for out_aval in _out_cts_avals)))
                vjp_polymorphic_shapes = [args_polymorphic_shapes, out_cts_polymorphic_shapes]
            out_cts = tree_util.tree_unflatten(out_tree_thunk(), out_cts_flat)
            with tf.name_scope('jax2tf_vjp'):
                in_cts = convert(fun_vjp_jax, with_gradient=False, polymorphic_shapes=vjp_polymorphic_shapes)(args, out_cts)
            return in_cts
        try:
            global _shape_env
            assert (not _shape_env), f'Unexpected shape environment {_shape_env}'
            global _enable_xla
            prev_enable_xla = _enable_xla
            _enable_xla = enable_xla
            _shape_env = shapeenv
            if with_gradient:

                @tf.custom_gradient
                def converted_fun_flat_with_custom_gradient(*args_flat: TfVal) -> TfVal:
                    out_with_avals = _interpret_fun(flat_fun, args_flat, args_avals_flat)
                    (outs, out_avals) = util.unzip2(out_with_avals)
                    return (tuple(outs), functools.partial(converted_grad_fn, _out_cts_avals=tuple(out_avals)))
                out_flat = converted_fun_flat_with_custom_gradient(*args_flat)
            else:
                out_flat_raw = _interpret_fun(flat_fun, args_flat, args_avals_flat)
                message = 'The jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients'
                out_flat = [tf.raw_ops.PreventGradient(input=o, message=message) for (o, _) in out_flat_raw]
        finally:
            _shape_env = {}
            _enable_xla = prev_enable_xla
        out_flat = [tf.identity(x, 'jax2tf_out') for x in out_flat]
        out = tree_util.tree_unflatten(out_tree_thunk(), out_flat)
        return out
    return converted_fun