original_docstring: Generate an operator to evaluate Jacobian-vector-products for a function.
  
      The passed function should take as arguments and return symbolic array(s),
      and likewise the returned operator will take symbolic array arguments.
  
      For a single argument function `func`, `n`-dimensional input array `x` and
      `n`-dimensional 'vector' array `v` of the same shape as `x` then we have
      the following equivalence
  
          jacobian_vector_product(func)(x)(v) == (
              tensordot(jacobian(func)(x), v, 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
              Jacobian-vector-product function and value of `func` as a 2-tuple
              (True) or just the Jacobian-vector-product function (False).
  
      Returns:
          Callable[..., Union[Callable[Array, Array], Tuple[Callable, Array]]]:
              Generated Jacobian-vector-product operator.
      
docstring: Generate an operator to evaluate Jacobian-vector-products for a function.
  The passed function should take as arguments and return symbolic array(s),
  and likewise the returned operator will take symbolic array arguments.
  
  
  
  
  
  
docstring_tokens: Generate an operator to evaluate Jacobian - vector - products for a function . The passed function should take as arguments and return symbolic array ( s ) and likewise the returned operator will take symbolic array arguments .
code: def jacobian_vector_product(func, wrt=0, return_aux=False):
      """Generate an operator to evaluate Jacobian-vector-products for a function.
  
      The passed function should take as arguments and return symbolic array(s),
      and likewise the returned operator will take symbolic array arguments.
  
      For a single argument function `func`, `n`-dimensional input array `x` and
      `n`-dimensional 'vector' array `v` of the same shape as `x` then we have
      the following equivalence
  
          jacobian_vector_product(func)(x)(v) == (
              tensordot(jacobian(func)(x), v, 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
              Jacobian-vector-product function and value of `func` as a 2-tuple
              (True) or just the Jacobian-vector-product function (False).
  
      Returns:
          Callable[..., Union[Callable[Array, Array], Tuple[Callable, Array]]]:
              Generated Jacobian-vector-product operator.
      """
  
      @_wrap_derived(func, 'jvp', 'Jacobian-vector-product')
      def jvp_func(*args):
          val = _get_sympy_func(func)(*args)
          jacob = _jacobian_transpose(
              sym.diff(val, args[wrt]), val.shape, args[wrt].shape)
          v = _named_array('v', args[wrt].shape)
          jvp = _generalised_dot(jacob, v, val.shape)
          v_jvp = FunctionExpression((v,), jvp)
          return (v_jvp, val) if return_aux else v_jvp
  
      return jvp_func
code_tokens: def jacobian_vector_product ( func , wrt = 0 , return_aux = False ) : @ _wrap_derived ( func , 'jvp' , 'Jacobian-vector-product' ) def jvp_func ( * args ) : val = _get_sympy_func ( func ) ( * args ) jacob = _jacobian_transpose ( sym . diff ( val , args [ wrt ] ) , val . shape , args [ wrt ] . shape ) v = _named_array ( 'v' , args [ wrt ] . shape ) jvp = _generalised_dot ( jacob , v , val . shape ) v_jvp = FunctionExpression ( ( v , ) , jvp ) return ( v_jvp , val ) if return_aux else v_jvp return jvp_func
short_docstring: Generate an operator to evaluate Jacobian-vector-products for a function.
short_docstring_tokens: Generate an operator to evaluate Jacobian - vector - products for a function .
comment: """Generate an operator to evaluate Jacobian-vector-products for a function.
  
      The passed function should take as arguments and return symbolic array(s),
      and likewise the returned operator will take symbolic array arguments.
  
      For a single argument function `func`, `n`-dimensional input array `x` and
      `n`-dimensional 'vector' array `v` of the same shape as `x` then we have
      the following equivalence
  
          jacobian_vector_product(func)(x)(v) == (
              tensordot(jacobian(func)(x), v, 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
              Jacobian-vector-product function and value of `func` as a 2-tuple
              (True) or just the Jacobian-vector-product function (False).
  
      Returns:
          Callable[..., Union[Callable[Array, Array], Tuple[Callable, Array]]]:
              Generated Jacobian-vector-product operator.
      """