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