(f, has_aux=False)
return _f
Returns a function that computes gradient of f.

  Gradients can only be computed through numpy and tensorflow operations and not
  through python float operations and values.

  Args:
    f: a function of type (params, *args) -> scalar. 'params' can be a nested
      structure (made of lists and tuples) of ndarrays and the gradient is
      evaluated against it. `scalar` is a scalar ndarray.
    has_aux: bool, indicates whether fun returns a pair where the first element
      is considered the output of the mathematical function to be differentiated
      and the second element is auxiliary data.

  Returns:
    A gradient function of type (params, *args) -> gradients, where the result
    'gradients' has the same structure and shapes as 'params'.
Returns a function that computes gradient of f.
Returns a function that computes gradient of f .
def grad(f, has_aux=False):
  """Returns a function that computes gradient of f.

  Gradients can only be computed through numpy and tensorflow operations and not
  through python float operations and values.

  Args:
    f: a function of type (params, *args) -> scalar. 'params' can be a nested
      structure (made of lists and tuples) of ndarrays and the gradient is
      evaluated against it. `scalar` is a scalar ndarray.
    has_aux: bool, indicates whether fun returns a pair where the first element
      is considered the output of the mathematical function to be differentiated
      and the second element is auxiliary data.

  Returns:
    A gradient function of type (params, *args) -> gradients, where the result
    'gradients' has the same structure and shapes as 'params'.
  """

  def check_loss_shape(np_loss):
    if not isinstance(np_loss, tf_np.ndarray):
      raise ValueError(
          "The result of the function to take gradient must be an ndarray.")
    if not np_loss.shape.is_compatible_with([]):
      raise ValueError(
          "The result of the function to take gradient must be a scalar.")

  def _f(params, *args):
    """The gradient function to be returned."""
    with tf.GradientTape() as g:
      g.watch(tf.nest.flatten(params))
      outputs = f(params, *args)
      if has_aux:
        np_loss, aux = outputs
      else:
        np_loss = outputs
      check_loss_shape(np_loss)
      tf_grads = g.gradient(np_loss, params)
      if has_aux:
        res = (tf_grads, aux)
      else:
        res = tf_grads
      return _tf_to_np(res)

  return _f
def grad ( f , has_aux = False ) : def check_loss_shape ( np_loss ) : if not isinstance ( np_loss , tf_np . ndarray ) : raise ValueError ( "The result of the function to take gradient must be an ndarray." ) if not np_loss . shape . is_compatible_with ( [ ] ) : raise ValueError ( "The result of the function to take gradient must be a scalar." ) def _f ( params , * args ) : """The gradient function to be returned.""" with tf . GradientTape ( ) as g : g . watch ( tf . nest . flatten ( params ) ) outputs = f ( params , * args ) if has_aux : np_loss , aux = outputs else : np_loss = outputs check_loss_shape ( np_loss ) tf_grads = g . gradient ( np_loss , params ) if has_aux : res = ( tf_grads , aux ) else : res = tf_grads return _tf_to_np ( res ) return _f