(f, *primals, has_aux=False)
return ret
Returns the result and the VJP function of `f`.

  This function returns the result and the vector-Jacobian-product (VJP)
  function of `f`.

  Args:
    f: a function from (nested structures of) tf_np.ndarrays to a (nested
      structure of) tf_np.ndarray. If `has_aux` is True, it should return an
      extra output.
    *primals: the inputs to be fed to `f`.
    has_aux: if True, the second output of `f` will be regarded as an auxiliary,
      non-differentiable output that will be ignored by the VJP function.

  Returns:
    A pair `(y, vjpfun)` if `has_aux` is False; a tuple `(y, vjpfun, aux)`
    otherwise. `y` and `aux` are the outputs of `f`, i.e. `y, aux =
    f(*primals)`. `vjpfun` is a function `dx = vjpfun(dy)`, where `dy` is the
    cotengents of `y`, having the same structures, shapes and dtypes as
    `y`. `dx` is the cotengents of `x`, having the same structures, shapes and
    dtypes as `x`.
Returns the result and the VJP function of `f`.
Returns the result and the VJP function of f .
def vjp(f, *primals, has_aux=False):
  """Returns the result and the VJP function of `f`.

  This function returns the result and the vector-Jacobian-product (VJP)
  function of `f`.

  Args:
    f: a function from (nested structures of) tf_np.ndarrays to a (nested
      structure of) tf_np.ndarray. If `has_aux` is True, it should return an
      extra output.
    *primals: the inputs to be fed to `f`.
    has_aux: if True, the second output of `f` will be regarded as an auxiliary,
      non-differentiable output that will be ignored by the VJP function.

  Returns:
    A pair `(y, vjpfun)` if `has_aux` is False; a tuple `(y, vjpfun, aux)`
    otherwise. `y` and `aux` are the outputs of `f`, i.e. `y, aux =
    f(*primals)`. `vjpfun` is a function `dx = vjpfun(dy)`, where `dy` is the
    cotengents of `y`, having the same structures, shapes and dtypes as
    `y`. `dx` is the cotengents of `x`, having the same structures, shapes and
    dtypes as `x`.
  """
  with tf.GradientTape(persistent=True) as tape:
    tape.watch(tf.nest.flatten(primals))
    outputs = f(*primals)
    if has_aux:
      np_out, aux = outputs
    else:
      np_out = outputs

    def _vjp(dy):
      tf_dx = tape.gradient(np_out, primals, output_gradients=dy)
      return _tf_to_np(tf_dx)

  if has_aux:
    ret = (np_out, _vjp, aux)
  else:
    ret = (np_out, _vjp)
  return ret
def vjp ( f , * primals , has_aux = False ) : with tf . GradientTape ( persistent = True ) as tape : tape . watch ( tf . nest . flatten ( primals ) ) outputs = f ( * primals ) if has_aux : np_out , aux = outputs else : np_out = outputs def _vjp ( dy ) : tf_dx = tape . gradient ( np_out , primals , output_gradients = dy ) return _tf_to_np ( tf_dx ) if has_aux : ret = ( np_out , _vjp , aux ) else : ret = ( np_out , _vjp ) return ret