original_docstring: Generate an operator to evaluate vector-Jacobian-products for a function.
  
      The passed function should take as arguments and return symbolic array(s),
      and likewise the returned function will act on symbolic arrays.
  
      For a single argument function `func`, input array `x` and `n`-dimensional
      'vector' array `v` of the same shape as `func(x)` then we have the
      following equivalence
  
          vector_jacobian_product(func)(x)(v) == (
              tensordot(v, jacobian(func)(x), n))
  
      where `tensordot` follows its NumPy semantics, i.e. `tensordot(a, b, n)`
      sums the products of components of `a` and `b` over the last `n` axes
      (dimensions) of `a` and first `n` dimensions of `b`.
  
      Args:
          func (Callable[..., Array]): Function which takes one or more arrays as
              arguments and returns an array.
          wrt (int): Index of argument to take derivatives with respect to.
          return_aux (bool): Whether the generated operator returns the
              vector-Jacobian-product function and value of `func` as a 2-tuple
              (True) or just the vector-Jacobian-product function (False).
  
      Returns:
          Callable[..., Union[Callable[Array, Array], Tuple[Callable, Array]]]:
              Generated vector-Jacobian-product operator.
      
docstring: Generate an operator to evaluate vector-Jacobian-products for a function.
  The passed function should take as arguments and return symbolic array(s),
  and likewise the returned function will act on symbolic arrays.
  
  
  
  
  
  
docstring_tokens: Generate an operator to evaluate vector - Jacobian - products for a function . The passed function should take as arguments and return symbolic array ( s ) and likewise the returned function will act on symbolic arrays .
code: def vector_jacobian_product(func, wrt=0, return_aux=False):
      """Generate an operator to evaluate vector-Jacobian-products for a function.
  
      The passed function should take as arguments and return symbolic array(s),
      and likewise the returned function will act on symbolic arrays.
  
      For a single argument function `func`, input array `x` and `n`-dimensional
      'vector' array `v` of the same shape as `func(x)` then we have the
      following equivalence
  
          vector_jacobian_product(func)(x)(v) == (
              tensordot(v, jacobian(func)(x), n))
  
      where `tensordot` follows its NumPy semantics, i.e. `tensordot(a, b, n)`
      sums the products of components of `a` and `b` over the last `n` axes
      (dimensions) of `a` and first `n` dimensions of `b`.
  
      Args:
          func (Callable[..., Array]): Function which takes one or more arrays as
              arguments and returns an array.
          wrt (int): Index of argument to take derivatives with respect to.
          return_aux (bool): Whether the generated operator returns the
              vector-Jacobian-product function and value of `func` as a 2-tuple
              (True) or just the vector-Jacobian-product function (False).
  
      Returns:
          Callable[..., Union[Callable[Array, Array], Tuple[Callable, Array]]]:
              Generated vector-Jacobian-product operator.
      """
  
      @_wrap_derived(func, 'vjp', 'vector-Jacobian-product')
      def vjp_func(*args):
          val = _get_sympy_func(func)(*args)
          jacob_transposed = sym.diff(val, args[wrt])
          v = _named_array('v', val.shape)
          vjp = _generalised_dot(jacob_transposed, v, args[wrt].shape)
          v_vjp = FunctionExpression((v,), vjp)
          return (v_vjp, val) if return_aux else v_vjp
  
      return vjp_func
code_tokens: def vector_jacobian_product ( func , wrt = 0 , return_aux = False ) : @ _wrap_derived ( func , 'vjp' , 'vector-Jacobian-product' ) def vjp_func ( * args ) : val = _get_sympy_func ( func ) ( * args ) jacob_transposed = sym . diff ( val , args [ wrt ] ) v = _named_array ( 'v' , val . shape ) vjp = _generalised_dot ( jacob_transposed , v , args [ wrt ] . shape ) v_vjp = FunctionExpression ( ( v , ) , vjp ) return ( v_vjp , val ) if return_aux else v_vjp return vjp_func
short_docstring: Generate an operator to evaluate vector-Jacobian-products for a function.
short_docstring_tokens: Generate an operator to evaluate vector - Jacobian - products for a function .
comment: """Generate an operator to evaluate vector-Jacobian-products for a function.
  
      The passed function should take as arguments and return symbolic array(s),
      and likewise the returned function will act on symbolic arrays.
  
      For a single argument function `func`, input array `x` and `n`-dimensional
      'vector' array `v` of the same shape as `func(x)` then we have the
      following equivalence
  
          vector_jacobian_product(func)(x)(v) == (
              tensordot(v, jacobian(func)(x), n))
  
      where `tensordot` follows its NumPy semantics, i.e. `tensordot(a, b, n)`
      sums the products of components of `a` and `b` over the last `n` axes
      (dimensions) of `a` and first `n` dimensions of `b`.
  
      Args:
          func (Callable[..., Array]): Function which takes one or more arrays as
              arguments and returns an array.
          wrt (int): Index of argument to take derivatives with respect to.
          return_aux (bool): Whether the generated operator returns the
              vector-Jacobian-product function and value of `func` as a 2-tuple
              (True) or just the vector-Jacobian-product function (False).
  
      Returns:
          Callable[..., Union[Callable[Array, Array], Tuple[Callable, Array]]]:
              Generated vector-Jacobian-product operator.
      """