
    [Th                      $   S SK r S SKJr  SSKJr  / SQrS rSS jrS r	S	 r
S
 rS rS r    SS jrS rSS jrSS jrS\\ R&                  S4   S\\S4   S\\ R&                  S4   4S jrSS jr    SS jr    SS jrSS jrSS jrg)     N)_vmap   )
forward_ad)vjpjvpjacobianhessianhvpvhpc                 v    [        U [        5      (       a  U $ [        U [        5      (       a  [        U 5      $ U 4$ N)
isinstancetuplelist)xs    Q/var/www/auris/envauris/lib/python3.13/site-packages/torch/autograd/functional.py_as_tuple_nocheckr      s1    !U	At		Qxt    c                 Z   Uc  Uc  [        U 5      $ Sn[        U [        5      (       d  U 4n Sn[        U 5       Hm  u  pE[        U[        R
                  5      (       a  M&  U(       a!  [        SU SU SU S[        U5       S3	5      e[        SU SU SU S[        U5       S3	5      e   X04$ )	NTFzThe z
 given to zF must be either a Tensor or a tuple of Tensors but the value at index z
 has type .z= must be either a Tensor or a tuple of Tensors but the given )r   r   r   	enumeratetorchTensor	TypeErrortype)inparg_namefn_nameis_inp_tupleiels         r   	_as_tupler"      s     GO %%Lc5!!f3"ell++8*Jwi 8''(cDH:Q@ 
  8*Jwi 8&Zz$r(1>    r   c                     [        U[        5      (       a>  [        U5      S:X  d   eUS   (       d  [        S U  5       5      n US   (       d  U S   n U $ U(       d  U S   n U $ )N   r   c              3   *   #    U  H	  oS    v   M     g7f)r   N ).0r!   s     r   	<genexpr>%_tuple_postprocess.<locals>.<genexpr>;   s     ,"1s   r   )r   r   len)res	to_unpacks     r   _tuple_postprocessr-   2   sd     )U##9~"""|,,,C|a&C J a&CJr   c                 b   / nU  H  nU(       ae  UR                   (       aT  UR                  (       d"  UR                  UR                  U5      5        MN  UR                  UR	                  5       5        Mo  UR                  UR                  5       R                  U5      5        M     [        U5      $ r   )requires_grad	is_sparseappendview_asclonedetachrequires_grad_r   )inputscreate_graph
need_graphr+   r   s        r   _grad_preprocessr9   D   sw     CC--==

3;;s+, 

399;'JJszz|22:>?  :r   c                    ^ [        U S   [        R                  5      (       a  T(       d  [        S U  5       5      $ U $ [        U4S jU  5       5      $ )Nr   c              3   @   #    U  H  oR                  5       v   M     g 7fr   )r4   r'   r   s     r   r(   $_grad_postprocess.<locals>.<genexpr>`   s     8#   c              3   <   >#    U  H  n[        UT5      v   M     g 7fr   )_grad_postprocess)r'   r   r7   s     r   r(   r=   d   s     LVc&sL99Vs   )r   r   r   r   )r6   r7   s    `r   r@   r@   [   sD     &)U\\**8888MLVLLLr   c           	         [        U5      [        U 5      :w  a6  U(       a$  [        S[        U5       S[        U 5       S35      e[        S5      e[        [        X5      5       Hi  u  nu  pEUR	                  5       UR	                  5       :w  d  M,  SnU(       a  SU S3n[        U SUR	                  5        SUR	                  5        S35      e   g )	Nz*v is a tuple of invalid length: should be z	 but got r   z+The given v should contain a single Tensor. zEntry z in zv has invalid size: should be )r*   RuntimeErrorr   zipsize)votheris_other_tupleidxel_vel_otherprepends          r   _validate_vrM   g   s     5zSV<SZL	RUVWRXQYYZ[  LMM!*3q=!9d99;(--/)G"3%t,)9(--/9J)TXT]T]T_S``ab  ":r   c                 @   U(       d  g US;  a  [        S5      e[        U 5       Hv  u  p4Uc  [        SU S35      eUR                  (       a  M*  US:X  a  [        SU S35      eUS:X  a  [        S	U S
35      eUS:X  a  [        SU S35      e[        SU S35      e   g )N)outputsgrad_inputsr   r	   z*Invalid input_type to _check_requires_gradAThe output of the user-provided function is independent of input %. This is not allowed in strict mode.r	   z@The hessian of the user-provided function with respect to input z is independent of the input. This is not allowed in strict mode. You should ensure that your function is thrice differentiable and that the hessian depends on the inputs.r   ziWhile computing the hessian, found that the jacobian of the user-provided function with respect to input z is independent of the input. This is not allowed in strict mode. You should ensure that your function is twice differentiable and that the jacobian depends on the inputs (this would be violated by a linear function for example).rP   z#The gradient with respect to input z` is independent of the inputs of the user-provided function. This is not allowed in strict mode.Output z of the user-provided function does not require gradients. The outputs must be computed in a differentiable manner from the input when running in strict mode.)rC   r   r/   )r6   
input_typestrictr    r   s        r   _check_requires_gradrV   |   s   JJGHHF#;STUSV W7 7     Y&"VWXVY::  z)"778c :CC  },"9! =S S 
 #aS !4 4 ; $r   c           
         [        U [        5      (       d   eUc  S[        U 5      -  n[        U[        5      (       d   e[        U 5      [        U5      :X  d   eSnSn[        X5       H'  u  pUc  M
  UR                  (       d  M  Xh4-  nXy4-  nM)     [        U5      S:X  a  S[        U5      -  $ [
        R                  R                  UUUSUUUS9$ )Nr   r&   r   T)allow_unusedr7   retain_graphis_grads_batched)r   r   r*   rD   r/   r   autogradgrad)
rO   r6   grad_outputsr7   rY   rZ   new_outputsnew_grad_outputsoutgrad_outs
             r   _autograd_gradrb      s     gu%%%%W-lE****w<3|,,,,,.K13W3?s0006!K+ 4
 ;1V$$~~""%%- # 
 	
r   c                    US;  a  [        SU S35      eSn[        U 5       H  u  pgUcn  U(       aN  US:X  a  [        SU S35      eUS:X  a  [        S	U S
35      eUS:X  a  [        SU S35      e[        SU S35      e[        R                  " X   5      nOCU(       a<  U(       a5  UR                  (       d$  SU;  a  [        SU S35      e[        SU S35      eXW4-  nM     U$ )N)back
back_trickdouble_backdouble_back_trickzInvalid stage argument 'z' to _fill_in_zerosr&   rd   rQ   rR   re   z?The gradient with respect to the input is independent of entry z in the grad_outputs when using the double backward trick to compute forward mode gradients. This is not allowed in strict mode.rf   CThe jacobian of the user-provided function is independent of input zBThe hessian of the user-provided function is independent of entry z in the grad_jacobian. This is not allowed in strict mode as it prevents from using the double backward trick to replace forward mode AD.double<. This is not allowed in strict mode when create_graph=True.zBThe hessian of the user-provided function is independent of input )rC   r   r   
zeros_liker/   )gradsrefsrU   r7   stager+   r    grads_is           r   _fill_in_zerosrp      sa   
 NN5eW<OPQQ$&C&
?F?&!!"#HJ  l*&YZ[Y\WW 
 m+&!!"#HJ 
 '!!" $33  &&tw/G,w/D/D5(&!!"#_a 
 '!!"#_a 
 	zS 'V Jr   c                    [         R                  " 5          [        USS5      u  pQ[        XSS9nU " U6 n[        USS5      u  pv[	        USUS9  Ub&  [        US	S5      u  p[        X#S
S9n[        X&U5        O1[        U5      S:w  d  US   R                  5       S:w  a  [        S5      eSSS5        U(       a  SO[         R                  " 5       n	[         R                  " U	5         [        WXUS9n
[        XXCS5      nSSS5        [        WU5      n[        WU5      n[        UW5      [        UW5      4$ ! , (       d  f       N= f! , (       d  f       NO= f)a
  Compute the dot product between a vector ``v`` and the Jacobian of the given function at the point given by the inputs.

Args:
    func (function): a Python function that takes Tensor inputs and returns
        a tuple of Tensors or a Tensor.
    inputs (tuple of Tensors or Tensor): inputs to the function ``func``.
    v (tuple of Tensors or Tensor): The vector for which the vector
        Jacobian product is computed.  Must be the same size as the output
        of ``func``. This argument is optional when the output of ``func``
        contains a single element and (if it is not provided) will be set
        as a Tensor containing a single ``1``.
    create_graph (bool, optional): If ``True``, both the output and result
        will be computed in a differentiable way. Note that when ``strict``
        is ``False``, the result can not require gradients or be
        disconnected from the inputs.  Defaults to ``False``.
    strict (bool, optional): If ``True``, an error will be raised when we
        detect that there exists an input such that all the outputs are
        independent of it. If ``False``, we return a Tensor of zeros as the
        vjp for said inputs, which is the expected mathematical value.
        Defaults to ``False``.

Returns:
    output (tuple): tuple with:
        func_output (tuple of Tensors or Tensor): output of ``func(inputs)``

        vjp (tuple of Tensors or Tensor): result of the dot product with
        the same shape as the inputs.

Example:

    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
    >>> def exp_reducer(x):
    ...     return x.exp().sum(dim=1)
    >>> inputs = torch.rand(4, 4)
    >>> v = torch.ones(4)
    >>> # xdoctest: +IGNORE_WANT("non-deterministic")
    >>> vjp(exp_reducer, inputs, v)
    (tensor([5.7817, 7.2458, 5.7830, 6.7782]),
     tensor([[1.4458, 1.3962, 1.3042, 1.6354],
            [2.1288, 1.0652, 1.5483, 2.5035],
            [2.2046, 1.1292, 1.1432, 1.3059],
            [1.3225, 1.6652, 1.7753, 2.0152]]))

    >>> vjp(exp_reducer, inputs, v, create_graph=True)
    (tensor([5.7817, 7.2458, 5.7830, 6.7782], grad_fn=<SumBackward1>),
     tensor([[1.4458, 1.3962, 1.3042, 1.6354],
            [2.1288, 1.0652, 1.5483, 2.5035],
            [2.2046, 1.1292, 1.1432, 1.3059],
            [1.3225, 1.6652, 1.7753, 2.0152]], grad_fn=<MulBackward0>))

    >>> def adder(x, y):
    ...     return 2 * x + 3 * y
    >>> inputs = (torch.rand(2), torch.rand(2))
    >>> v = torch.ones(2)
    >>> vjp(adder, inputs, v)
    (tensor([2.4225, 2.3340]),
     (tensor([2., 2.]), tensor([3., 3.])))
r6   r   Tr7   r8   %outputs of the user-provided functionrO   rU   NrF   Fr   r   zjThe vector v can only be None if the user-provided function returns a single Tensor with a single element.r7   rd   )r   enable_gradr"   r9   rV   rM   r*   nelementrC   is_grad_enabledset_grad_enabledrb   rp   r@   r-   )funcr6   rF   r7   rU   is_inputs_tuplerO   is_outputs_tuple_rv   grad_resr   s               r   r   r     sP   v 
			"+FHe"D!&PTU-$-<e%
! 	Wi?=QU+DA %PA$457|q GAJ$7$7$9Q$>"= ! 
, '$E,A,A,CK				,!'6<PXvVL 
-
  6G
C
.Cg'78:L_;  ? 
	. 
-	,s   BD6%E6
E
Ec                    [         R                  " 5          [        USS5      u  pQ[        XSS9nUb&  [        USS5      u  pb[        X#SS9n[	        X!U5        O1[        U5      S:w  d  US	   R                  5       S:w  a  [        S
5      eU " U6 n[        USS5      u  p[        USUS9  [        S U 5       5      n	[        XqU	SS9n
[        U
SUS9  SSS5        U(       a8  [         R                  " 5          [        W
W	X#S9n[        UWXCS5      nSSS5        O[        W
W	X#S9n[        UWXCS5      n[        WU5      n[        WU5      n[        UW5      [        X5      4$ ! , (       d  f       N= f! , (       d  f       NN= f)aL
  Compute the dot product between the Jacobian of the given function at the point given by the inputs and a vector ``v``.

Args:
    func (function): a Python function that takes Tensor inputs and returns
        a tuple of Tensors or a Tensor.
    inputs (tuple of Tensors or Tensor): inputs to the function ``func``.
    v (tuple of Tensors or Tensor): The vector for which the Jacobian
        vector product is computed. Must be the same size as the input of
        ``func``. This argument is optional when the input to ``func``
        contains a single element and (if it is not provided) will be set
        as a Tensor containing a single ``1``.
    create_graph (bool, optional): If ``True``, both the output and result
        will be computed in a differentiable way. Note that when ``strict``
        is ``False``, the result can not require gradients or be
        disconnected from the inputs.  Defaults to ``False``.
    strict (bool, optional): If ``True``, an error will be raised when we
        detect that there exists an input such that all the outputs are
        independent of it. If ``False``, we return a Tensor of zeros as the
        jvp for said inputs, which is the expected mathematical value.
        Defaults to ``False``.

Returns:
    output (tuple): tuple with:
        func_output (tuple of Tensors or Tensor): output of ``func(inputs)``

        jvp (tuple of Tensors or Tensor): result of the dot product with
        the same shape as the output.

Note:
    ``autograd.functional.jvp`` computes the jvp by using the backward of
    the backward (sometimes called the double backwards trick). This is not
    the most performant way of computing the jvp. Please consider using
    :func:`torch.func.jvp` or the
    :ref:`low-level forward-mode AD API <forward-mode-ad>` instead.

Example:

    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
    >>> def exp_reducer(x):
    ...     return x.exp().sum(dim=1)
    >>> inputs = torch.rand(4, 4)
    >>> v = torch.ones(4, 4)
    >>> # xdoctest: +IGNORE_WANT("non-deterministic")
    >>> jvp(exp_reducer, inputs, v)
    (tensor([6.3090, 4.6742, 7.9114, 8.2106]),
     tensor([6.3090, 4.6742, 7.9114, 8.2106]))

    >>> jvp(exp_reducer, inputs, v, create_graph=True)
    (tensor([6.3090, 4.6742, 7.9114, 8.2106], grad_fn=<SumBackward1>),
     tensor([6.3090, 4.6742, 7.9114, 8.2106], grad_fn=<SqueezeBackward1>))

    >>> def adder(x, y):
    ...     return 2 * x + 3 * y
    >>> inputs = (torch.rand(2), torch.rand(2))
    >>> v = (torch.ones(2), torch.ones(2))
    >>> jvp(adder, inputs, v)
    (tensor([2.2399, 2.5005]),
     tensor([5., 5.]))

r6   r   Trr   NrF   Fr   r   rThe vector v can only be None if the input to the user-provided function is a single Tensor with a single element.rs   rO   rt   c              3   L   #    U  H  n[         R                  " US S9v   M     g7fT)r/   Nr   rk   )r'   r`   s     r   r(   jvp.<locals>.<genexpr>  s       
AH#ES5   "$ru   rP   re   )r   rv   r"   r9   rM   r*   rw   rC   rV   r   rb   rp   r@   r-   )rz   r6   rF   r7   rU   r{   r}   rO   r|   r]   rP   r~   r   s                r   r   r   f  s   z 
			"+FHe"D!&PTU=QU+DA %PA?36{a6!9#5#5#71#<"-  -$-<e%
! 	Wi?  
AH
 
 %WlQUV[-G; 
>  %\1H !7F,WC	 !  "q
 XwlS  6G
C
.Cg'78:L;  _ 
	@ ! s   B;E 5E1 
E.1
E?tensors.tensor_numelsreturnc                 $  ^ [        U 5      [        U5      :X  d   e[        U 5      S:  d   e[        U5      m[        U4S j[        X5       5       5      nSn[        X!5       H)  u  pEUR	                  U5      R                  S5        X5-  nM+     U$ )Nr   c              3   L   >#    U  H  u  pUR                  TU5      v   M     g 7fr   )	new_zeros)r'   tensortensor_numeltotal_numels      r   r(   0_construct_standard_basis_for.<locals>.<genexpr>  s+      $? F 	l33$?s   !$r   )r*   sumr   rD   diagonalfill_)r   r   chunksdiag_start_idxchunknumelr   s         @r   _construct_standard_basis_forr     s    * w<3}----w<!m$K $'$? F NF2~&,,Q/ 3 Mr   c           	      `  ^ ^^ U(       a  [        S5      e[        TSS5      u  nm/ mU(       a  [        S T 5       5      n[        TU5      nU UU4S jn[	        U5      " U5      nTu  p/ n[        X5       H  u  p/ n[        UR                  USS9T5       Hg  u  nnUR                  " / [        SUR                  5      QSP76 R                  / UR                  QUR                  Q75      nUR                  U5        Mi     UR                  U5        M     [        XU45      $ [        S	5      e)
Nztorch.autograd.functional.jacobian: `strict=True` and `strategy="forward-mode"` are not supported together (yet). Please either set `strict=False` or `strategy="reverse-mode"`.r6   r   c              3   @   #    U  H  oR                  5       v   M     g 7fr   r   )r'   inputs     r   r(   _jacfwd.<locals>.<genexpr>  s     ?u[[]]r>   c                   > [         R                  " 5          [        S [        T
U 5       5       5      n[	        T	" U6 S5      u  p#TR                  U5        / n/ nU Hg  n[         R                  " U5      u  pxUR                  U5        Ub  UR                  U5        MB  UR                  [        R                  " U5      5        Mi     TR                  U5        [        U5      sS S S 5        $ ! , (       d  f       g = f)Nc              3   p   #    U  H,  u  p[         R                  " XR                  U5      5      v   M.     g 7fr   )fwAD	make_dualr2   )r'   r   tangents      r   r(   '_jacfwd.<locals>.jvp.<locals>.<genexpr>  s.      $*? NN5//%*@AA*?s   46rO   )	r   
dual_levelr   rD   r"   r1   unpack_dualr   rk   )tangentsdual_inputs_is_outputs_tupledual_outputsjvprimal_outsdual_outprimalr   rz   r6   output_infos            r   r   _jacfwd.<locals>.jvp  s    "# $*-fh*?$  3<+&	3/! ""#45  ,H&*&6&6x&@OF&&v.*		'*		%"2"26":; !- "";/Ry' #""s   CC,,
C:r   dimr   zlComputing Jacobian using forward-AD or forward-over-reverse Hessian isonly implemented for `vectorize=True`.)rC   r"   r   r   r   rD   splitpermuterangendimreshapeshaper1   r-   NotImplementedError)rz   r6   rU   	vectorizer{   input_numelsr   r   outputs_before_splitr|   rO   jacobian_input_outputjac_output_ioutput_ijacobian_output_i_outputjacinput_jjacobian_input_i_output_jr   s   ``                @r   _jacfwdr     sI   )
 	
 (*EOVK??? 1F	!,  %Sz(3$/! "&)*>&H"L')$ #L$6$6|$6$KV TW -0KK,Oq#((9K,OQ,O,W,W5hnn5w}}5-) )//0IJ !U "(()AB 'I "!o#F
 	
 "5
 	
r   c           
      2  ^^^^^  US;   d   S5       eUS:X  a  T(       a  [        S5      e[        U TX45      $ [        R                  " 5          [	        TSS5      u  nm[        TTSS9mU " T6 n[	        US	S5      u  p[        US
US9  U(       Ga  U(       a  [        S5      e[        S U 5       5      m [        UT 5      n	[        S U 5       5      mUUUU 4S jn
U
" U	5      n/ n[        UT5       Hu  u  p/ n[        UR                  T SS9U5       H?  u  nnUR                  UR                  UR                  -   5      nUR                  U5        MA     UR                  U5        Mw     [        [        U6 5      n[        UT5      n[!        UX45      sSSS5        $ Sn[#        U5       GH7  u  nm[        S [%        ['        T5      5       5       5      n[%        TR)                  5       5       H  n[+        TR-                  S5      U   4TSTS9n[#        [        UUT5      5       H  u  nu  nnnUbC  U(       a)  T(       a"  UR.                  (       d  SU S3n[        U5      eUR                  U5        MP  U(       a  SU SU S3n[        U5      eUR                  [        R0                  " U5      5        M     M     U[        UU4S j[#        U5       5       5      4-  nGM:     [        UT5      n[!        UX45      sSSS5        $ ! , (       d  f       g= f)a9  Compute the Jacobian of a given function.

Args:
    func (function): a Python function that takes Tensor inputs and returns
        a tuple of Tensors or a Tensor.
    inputs (tuple of Tensors or Tensor): inputs to the function ``func``.
    create_graph (bool, optional): If ``True``, the Jacobian will be
        computed in a differentiable manner. Note that when ``strict`` is
        ``False``, the result can not require gradients or be disconnected
        from the inputs.  Defaults to ``False``.
    strict (bool, optional): If ``True``, an error will be raised when we
        detect that there exists an input such that all the outputs are
        independent of it. If ``False``, we return a Tensor of zeros as the
        jacobian for said inputs, which is the expected mathematical value.
        Defaults to ``False``.
    vectorize (bool, optional): This feature is experimental.
        Please consider using :func:`torch.func.jacrev` or
        :func:`torch.func.jacfwd` instead if you are looking for something
        less experimental and more performant.
        When computing the jacobian, usually we invoke
        ``autograd.grad`` once per row of the jacobian. If this flag is
        ``True``, we perform only a single ``autograd.grad`` call with
        ``batched_grad=True`` which uses the vmap prototype feature.
        Though this should lead to performance improvements in many cases,
        because this feature is still experimental, there may be performance
        cliffs. See :func:`torch.autograd.grad`'s ``batched_grad`` parameter for
        more information.
    strategy (str, optional): Set to ``"forward-mode"`` or ``"reverse-mode"`` to
        determine whether the Jacobian will be computed with forward or reverse
        mode AD. Currently, ``"forward-mode"`` requires ``vectorized=True``.
        Defaults to ``"reverse-mode"``. If ``func`` has more outputs than
        inputs, ``"forward-mode"`` tends to be more performant. Otherwise,
        prefer to use ``"reverse-mode"``.

Returns:
    Jacobian (Tensor or nested tuple of Tensors): if there is a single
    input and output, this will be a single Tensor containing the
    Jacobian for the linearized inputs and output. If one of the two is
    a tuple, then the Jacobian will be a tuple of Tensors. If both of
    them are tuples, then the Jacobian will be a tuple of tuple of
    Tensors where ``Jacobian[i][j]`` will contain the Jacobian of the
    ``i``\th output and ``j``\th input and will have as size the
    concatenation of the sizes of the corresponding output and the
    corresponding input and will have same dtype and device as the
    corresponding input. If strategy is ``forward-mode``, the dtype will be
    that of the output; otherwise, the input.

Example:

    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
    >>> def exp_reducer(x):
    ...     return x.exp().sum(dim=1)
    >>> inputs = torch.rand(2, 2)
    >>> # xdoctest: +IGNORE_WANT("non-deterministic")
    >>> jacobian(exp_reducer, inputs)
    tensor([[[1.4917, 2.4352],
             [0.0000, 0.0000]],
            [[0.0000, 0.0000],
             [2.4369, 2.3799]]])

    >>> jacobian(exp_reducer, inputs, create_graph=True)
    tensor([[[1.4917, 2.4352],
             [0.0000, 0.0000]],
            [[0.0000, 0.0000],
             [2.4369, 2.3799]]], grad_fn=<ViewBackward>)

    >>> def exp_adder(x, y):
    ...     return 2 * x.exp() + 3 * y
    >>> inputs = (torch.rand(2), torch.rand(2))
    >>> jacobian(exp_adder, inputs)
    (tensor([[2.8052, 0.0000],
            [0.0000, 3.3963]]),
     tensor([[3., 0.],
             [0., 3.]]))
forward-modereverse-modezExpected strategy to be either "forward-mode" or "reverse-mode". Hint: If your function has more outputs than inputs, "forward-mode" tends to be more performant. Otherwise, prefer to use "reverse-mode".r   ztorch.autograd.functional.jacobian: `create_graph=True` and `strategy="forward-mode"` are not supported together (yet). Please either set `create_graph=False` or `strategy="reverse-mode"`.r6   r   Trr   rs   rO   rt   ztorch.autograd.functional.jacobian: `strict=True` and `vectorized=True` are not supported together. Please either set `strict=False` or `vectorize=False`.c              3   @   #    U  H  oR                  5       v   M     g 7fr   r   r'   outputs     r   r(   jacobian.<locals>.<genexpr>  s     !GwV,,..wr>   c              3   B   #    U  H  oR                  S 5      v   M     g7f)N)r   r   s     r   r(   r     s      J'!3!3'   c           
         > [        [        TTU TSS95      n[        U5       HN  u  p#Ub  M
  [        R                  " TU   5      R                  [        T5      4TU   R                  -   5      X'   MP     [        U5      $ )NT)r7   rZ   )	r   rb   r   r   rk   expandr   r   r   )grad_outputvjel_idxvj_elr7   flat_outputsr6   output_numelss       r   r   jacobian.<locals>.vjp  s    "$#%1)- &/r]MF( !&!1!1&.!A!H!H]+-v0D0DD"BJ &3 Ry r   r   r   Nr&   c              3   &   #    U  H  n/ v   M	     g 7fr   r&   )r'   r}   s     r   r(   r     s     4TASARASs   r   )rY   r7   rh   rj   rS   z7 of the user-provided function is independent of input rR   c              3      >#    U  HM  u  p[         R                  " US S9R                  TR                  5       TU   R                  5       -   5      v   MO     g7f)r   r   N)r   stackviewrE   )r'   r   jac_i_elr6   r`   s      r   r(   r   2  sU       /?* KKa055
VF^%8%8%::  /?s   AA)r   r   r   rv   r"   r9   rV   rC   r   r   rD   r   r   r   r1   r@   r-   r   r   r*   rw   rb   r   r/   rk   )!rz   r6   r7   rU   r   strategyr{   rO   r|   r]   r   jacobians_of_flat_outputr   jac_input_iinput_ijacobian_input_i_outputr   output_jr   jacobian_output_inputr   r    jac_ijr   r   r   r   inp_elmsgr   r`   r   s!    ``                           @@@r   r   r   >  sU   f 77 	37
 >!%-  tVV77				"+FHj"I!&|PTU-$-<j%
! 	Wi?") ^ "!Gw!GGM8-PL  J' JJL! !$ (+<'8$ %'!(+,Df(M$*,'%(%%m%;W&MC 14'--9W0X-+223LM	&
 &,,-DE )N %*#/D*E$F!$5%|%! &%(8'JU 
	\ .0(FAs/44Ts6{AS4T/TE3<<>*#[[_Q')!%!-	 :Cr6*:5F5Xuf (!l5;N;N!889s ;F!F  
 #/s"33 .!")! -88>x @/!/  
 #/s"33 (8(8(@A): +<   /8.>	  HC )T %X|<!(-=,OPy 
		s   ELE!L
Lc           	         ^ ^^^	 [        USS5      u  paTS;   d   S5       eU U4S jm	U	UU4S jn[        UUUTUTS9n[        XU45      $ )as  Compute the Hessian of a given scalar function.

Args:
    func (function): a Python function that takes Tensor inputs and returns
        a Tensor with a single element.
    inputs (tuple of Tensors or Tensor): inputs to the function ``func``.
    create_graph (bool, optional): If ``True``, the Hessian will be computed in
        a differentiable manner. Note that when ``strict`` is ``False``, the result can not
        require gradients or be disconnected from the inputs.
        Defaults to ``False``.
    strict (bool, optional): If ``True``, an error will be raised when we detect that there exists an input
        such that all the outputs are independent of it. If ``False``, we return a Tensor of zeros as the
        hessian for said inputs, which is the expected mathematical value.
        Defaults to ``False``.
    vectorize (bool, optional): This feature is experimental.
        Please consider using :func:`torch.func.hessian`
        instead if you are looking for something less experimental and more performant.
        When computing the hessian, usually we invoke
        ``autograd.grad`` once per row of the hessian. If this flag is
        ``True``, we use the vmap prototype feature as the backend to
        vectorize calls to ``autograd.grad`` so we only invoke it once
        instead of once per row. This should lead to performance
        improvements in many use cases, however, due to this feature
        being incomplete, there may be performance cliffs. Please
        use `torch._C._debug_only_display_vmap_fallback_warnings(True)`
        to show any performance warnings and file us issues if
        warnings exist for your use case. Defaults to ``False``.
    outer_jacobian_strategy (str, optional): The Hessian is computed by
        computing the Jacobian of a Jacobian. The inner Jacobian is always
        computed in reverse-mode AD. Setting strategy to ``"forward-mode"``
        or ``"reverse-mode"`` determines whether the outer Jacobian will be
        computed with forward or reverse mode AD. Currently, computing the outer
        Jacobian in ``"forward-mode"`` requires ``vectorized=True``. Defaults
        to ``"reverse-mode"``.

Returns:
    Hessian (Tensor or a tuple of tuple of Tensors): if there is a single input,
    this will be a single Tensor containing the Hessian for the input.
    If it is a tuple, then the Hessian will be a tuple of tuples where
    ``Hessian[i][j]`` will contain the Hessian of the ``i``\th input
    and ``j``\th input with size the sum of the size of the ``i``\th input plus
    the size of the ``j``\th input. ``Hessian[i][j]`` will have the same
    dtype and device as the corresponding ``i``\th input.

Example:

    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
    >>> def pow_reducer(x):
    ...     return x.pow(3).sum()
    >>> inputs = torch.rand(2, 2)
    >>> # xdoctest: +IGNORE_WANT("non-deterministic")
    >>> hessian(pow_reducer, inputs)
    tensor([[[[5.2265, 0.0000],
              [0.0000, 0.0000]],
             [[0.0000, 4.8221],
              [0.0000, 0.0000]]],
            [[[0.0000, 0.0000],
              [1.9456, 0.0000]],
             [[0.0000, 0.0000],
              [0.0000, 3.2550]]]])

    >>> hessian(pow_reducer, inputs, create_graph=True)
    tensor([[[[5.2265, 0.0000],
              [0.0000, 0.0000]],
             [[0.0000, 4.8221],
              [0.0000, 0.0000]]],
            [[[0.0000, 0.0000],
              [1.9456, 0.0000]],
             [[0.0000, 0.0000],
              [0.0000, 3.2550]]]], grad_fn=<ViewBackward>)


    >>> def pow_adder_reducer(x, y):
    ...     return (2 * x.pow(2) + 3 * y.pow(2)).sum()
    >>> inputs = (torch.rand(2), torch.rand(2))
    >>> hessian(pow_adder_reducer, inputs)
    ((tensor([[4., 0.],
              [0., 4.]]),
      tensor([[0., 0.],
              [0., 0.]])),
     (tensor([[0., 0.],
              [0., 0.]]),
      tensor([[6., 0.],
              [0., 6.]])))
r6   r	   r   z@Expected strategy to be either "forward-mode" or "reverse-mode".c                    > T" U 6 n[        USS5      u  p#[        USTS9  U(       d  [        U[        R                  5      (       d  [        S5      eUR                  5       S:w  a  [        S5      eUR                  5       $ )Nrs   r	   rO   rt   z;The function given to hessian should return a single Tensorr   zTThe Tensor returned by the function given to hessian should contain a single element)r"   rV   r   r   r   rC   rw   squeeze)r   r`   is_out_tuplet_outrz   rU   s       r   ensure_single_output_function.hessian.<locals>.ensure_single_output_function  s    Cj'8)
 	UIf=z#u||<<M  <<>Qf  {{}r   c                  d   > TS:X  a  [        S U  5       5      n [        TU SS9n[        USTS9  U$ )Nr   c              3   B   #    U  H  oR                  S 5      v   M     g7f)TN)r5   )r'   ts     r   r(   ,hessian.<locals>.jac_func.<locals>.<genexpr>  s     <1((..r   Tru   r   rt   )r   r   rV   )r   r   r   outer_jacobian_strategyrU   s     r   jac_funchessian.<locals>.jac_func  s=    "n4 <<<C4cMS*V<
r   )r7   rU   r   r   )r"   r   r-   )
rz   r6   r7   rU   r   r   r{   r   r+   r   s
   `  ` `   @r   r	   r	   ?  su    z ()DO" '  J JJ 
& !(C c_#EFFr   c                    [         R                  " 5          [        USS5      u  pQ[        XSS9nUb&  [        USS5      u  pb[        X#SS9n[	        X!U5        O1[        U5      S:w  d  US	   R                  5       S:w  a  [        S
5      eU " U6 n[        USS5      u  p[        USUS9  U(       d"  [        US	   [         R                  5      (       d  [        S5      eUS	   R                  5       S:w  a  [        S5      e[        XqSS9n	[        U	SUS9  SSS5        U(       a  SO[         R                  " 5       n
[         R                  " U
5         [        W	XUS9n[        XXCS5      nSSS5        [        WU5      n[        WU5      n[!        UW5      [!        UW5      4$ ! , (       d  f       N= f! , (       d  f       NO= f)a	  Compute the dot product between vector ``v`` and Hessian of a  given scalar function at a specified point.

Args:
    func (function): a Python function that takes Tensor inputs and returns
        a Tensor with a single element.
    inputs (tuple of Tensors or Tensor): inputs to the function ``func``.
    v (tuple of Tensors or Tensor): The vector for which the vector Hessian
        product is computed. Must be the same size as the input of
        ``func``. This argument is optional when ``func``'s input contains
        a single element and (if it is not provided) will be set as a
        Tensor containing a single ``1``.
    create_graph (bool, optional): If ``True``, both the output and result
        will be computed in a differentiable way. Note that when ``strict``
        is ``False``, the result can not require gradients or be
        disconnected from the inputs.
        Defaults to ``False``.
    strict (bool, optional): If ``True``, an error will be raised when we
        detect that there exists an input such that all the outputs are
        independent of it. If ``False``, we return a Tensor of zeros as the
        vhp for said inputs, which is the expected mathematical value.
        Defaults to ``False``.

Returns:
    output (tuple): tuple with:
        func_output (tuple of Tensors or Tensor): output of ``func(inputs)``

        vhp (tuple of Tensors or Tensor): result of the dot product with the
        same shape as the inputs.

Example:

    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
    >>> def pow_reducer(x):
    ...     return x.pow(3).sum()
    >>> inputs = torch.rand(2, 2)
    >>> v = torch.ones(2, 2)
    >>> # xdoctest: +IGNORE_WANT("non-deterministic")
    >>> vhp(pow_reducer, inputs, v)
    (tensor(0.5591),
     tensor([[1.0689, 1.2431],
             [3.0989, 4.4456]]))
    >>> vhp(pow_reducer, inputs, v, create_graph=True)
    (tensor(0.5591, grad_fn=<SumBackward0>),
     tensor([[1.0689, 1.2431],
             [3.0989, 4.4456]], grad_fn=<MulBackward0>))
    >>> def pow_adder_reducer(x, y):
    ...     return (2 * x.pow(2) + 3 * y.pow(2)).sum()
    >>> inputs = (torch.rand(2), torch.rand(2))
    >>> v = (torch.zeros(2), torch.ones(2))
    >>> vhp(pow_adder_reducer, inputs, v)
    (tensor(4.8053),
     (tensor([0., 0.]),
      tensor([6., 6.])))
r6   r   Trr   NrF   Fr   r   r   rs   rO   rt   z7The function given to vhp should return a single TensorzPThe Tensor returned by the function given to vhp should contain a single elementru   r   rf   )r   rv   r"   r9   rM   r*   rw   rC   rV   r   r   rb   rx   ry   rp   r@   r-   )rz   r6   rF   r7   rU   r{   r}   rO   r|   r   rv   r~   r   s                r   r   r     s   n 
			"+FHe"D!&PTU=QU+DA %PA?36{a6!9#5#5#71#<"@  -$-<e%
! 	Wi?:gaj%,,#G#GI  1: A%b  W4@S*V<? 
B '$E,A,A,CK				,!#v|LXv]S 
-  6G
C
.Cg'78:L_;  S 
	D 
-	,s   C>F!F2!
F/2
G c                    [         R                  " 5          [        USS5      u  pQ[        XSS9nUb&  [        USS5      u  pb[        X#SS9n[	        X!U5        O1[        U5      S:w  d  US	   R                  5       S:w  a  [        S
5      eU " U6 n[        USS5      u  p[        USUS9  U(       d"  [        US	   [         R                  5      (       d  [        S5      eUS	   R                  5       S:w  a  [        S5      e[        XqSS9n	[        U	SUS9  [        S U 5       5      n
[        XU
SS9n[        U	SUS9  SSS5        U(       a  SO[         R                  " 5       n[         R                  " U5         [        WW
X#S9n[        XXCS5      nSSS5        [!        WU5      n[!        WU5      n[#        UW5      [#        UW5      4$ ! , (       d  f       N= f! , (       d  f       NO= f)aU
  Compute the dot product between the scalar function's Hessian and a vector ``v`` at a specified point.

Args:
    func (function): a Python function that takes Tensor inputs and returns
        a Tensor with a single element.
    inputs (tuple of Tensors or Tensor): inputs to the function ``func``.
    v (tuple of Tensors or Tensor): The vector for which the Hessian vector
        product is computed. Must be the same size as the input of
        ``func``. This argument is optional when ``func``'s input contains
        a single element and (if it is not provided) will be set as a
        Tensor containing a single ``1``.
    create_graph (bool, optional): If ``True``, both the output and result will be
        computed in a differentiable way. Note that when ``strict`` is
        ``False``, the result can not require gradients or be disconnected
        from the inputs.  Defaults to ``False``.
    strict (bool, optional): If ``True``, an error will be raised when we
        detect that there exists an input such that all the outputs are
        independent of it. If ``False``, we return a Tensor of zeros as the
        hvp for said inputs, which is the expected mathematical value.
        Defaults to ``False``.
Returns:
    output (tuple): tuple with:
        func_output (tuple of Tensors or Tensor): output of ``func(inputs)``

        hvp (tuple of Tensors or Tensor): result of the dot product with
        the same shape as the inputs.

Example:

    >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
    >>> def pow_reducer(x):
    ...     return x.pow(3).sum()
    >>> inputs = torch.rand(2, 2)
    >>> v = torch.ones(2, 2)
    >>> # xdoctest: +IGNORE_WANT("non-deterministic")
    >>> hvp(pow_reducer, inputs, v)
    (tensor(0.1448),
     tensor([[2.0239, 1.6456],
             [2.4988, 1.4310]]))

    >>> hvp(pow_reducer, inputs, v, create_graph=True)
    (tensor(0.1448, grad_fn=<SumBackward0>),
     tensor([[2.0239, 1.6456],
             [2.4988, 1.4310]], grad_fn=<MulBackward0>))


    >>> def pow_adder_reducer(x, y):
    ...     return (2 * x.pow(2) + 3 * y.pow(2)).sum()
    >>> inputs = (torch.rand(2), torch.rand(2))
    >>> v = (torch.zeros(2), torch.ones(2))
    >>> hvp(pow_adder_reducer, inputs, v)
    (tensor(2.3030),
     (tensor([0., 0.]),
      tensor([6., 6.])))

Note:

    This function is significantly slower than `vhp` due to backward mode AD constraints.
    If your functions is twice continuously differentiable, then hvp = vhp.t(). So if you
    know that your function satisfies this condition, you should use vhp instead that is
    much faster with the current implementation.

r6   r
   Trr   NrF   Fr   r   r   rs   rO   rt   z7The function given to hvp should return a single TensorzPThe Tensor returned by the function given to hvp should contain a single elementru   r   c              3   L   #    U  H  n[         R                  " US S9v   M     g7fr   r   r<   s     r   r(   hvp.<locals>.<genexpr>  s     Ufs))#TBfr   r	   rg   )r   rv   r"   r9   rM   r*   rw   rC   rV   r   r   rb   r   rx   ry   rp   r@   r-   )rz   r6   rF   r7   rU   r{   r}   rO   r|   r   grad_jacrf   rv   r~   r
   s                  r   r
   r
   .  s   @ 
			"+FHe"D!&PTU=QU+DA %PA?36{a6!9#5#5#71#<"@  -$-<e%
! 	Wi?:gaj%,,#G#GI  1: A%b  W4@S*V<UfUU$S(NS)F;I 
L '$E,A,A,CK				,!+xVf4G
 
-  6G
C
.Cg'78:L_;  a 
	N 
-	,s   D&G	8G	
G
G()NN)NFNF)NFF)FF)FFFr   )r   torch._vmap_internalsr   rB   r   r   __all__r   r"   r-   r9   r@   rM   rV   rb   rp   r   r   r   r   intr   r   r   r	   r   r
   r&   r   r   <module>r      s     '   >
6$.	M*(\ #
L4t\~nb 5<<$% 6;CHo 
5<< FA
N ~QH *GGTbJrr   