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
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