def custom_grad(f_vjp, f_original=None):
  """Decorator to define a function with a custom gradient.

  This function is very similar to `tf.custom_gradient`. See the documentation
  of `tf.custom_gradient` for detailed usage.

  The differences with `tf.custom_gradient` are:

  - All arguments and results are tf_np.ndarrays instead of tensors.

  - The `grad_fn` returned by `f_vjp` accepts and returns nested structures,
    unlike that in `tf.custom_gradient` which only accepts and returns lists.

  Args:
    f_vjp: the same as the `f` argument of `tf.custom_gradient`. Note that all
      inputs and outputs of `f_vjp` and of the `grad_fn` function it returns can
      be nested structures.
    f_original: (optional) not used.

  Returns:
    The same as `tf.custom_gradient`.
  """
  del f_original

  @tf.custom_gradient
  def tf_f(*tf_args, **tf_kwargs):
    np_args = _tf_to_np(tf_args)
    np_kwargs = _tf_to_np(tf_kwargs)
    np_y, np_vjp = f_vjp(*np_args, **np_kwargs)
    tf_y = np_y

    def tf_vjp(*flat_tf_dy):
      tf_dy = tf.nest.pack_sequence_as(tf_y, flat_tf_dy)
      np_dy = _tf_to_np(tf_dy)
      np_dx = np_vjp(np_dy)
      return tf.nest.flatten(np_dx)

    return tf_y, tf_vjp

  def np_f(*args, **kwargs):
    return _tf_to_np(tf_f(*args), **kwargs)

  return np_f
Decorator to define a function with a custom gradient.

  This function is very similar to `tf.custom_gradient`. See the documentation
  of `tf.custom_gradient` for detailed usage.

  The differences with `tf.custom_gradient` are:

  - All arguments and results are tf_np.ndarrays instead of tensors.

  - The `grad_fn` returned by `f_vjp` accepts and returns nested structures,
    unlike that in `tf.custom_gradient` which only accepts and returns lists.

  Args:
    f_vjp: the same as the `f` argument of `tf.custom_gradient`. Note that all
      inputs and outputs of `f_vjp` and of the `grad_fn` function it returns can
      be nested structures.
    f_original: (optional) not used.

  Returns:
    The same as `tf.custom_gradient`.
  
trax/tf_numpy/extensions/extensions.py