
    AVhR%                        d Z ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ ddlm	Z	 dd	lm
Z
 dd
l d Z  ej                  d      e         ej                  d      e        ej                  d      d        Z ej                  d      d        Z ej                  d      d        Z  ej                  d      e        ej                  d      d        Z ej                  d      d        Z ej                  d      d        Z ej                  d      d        Z ej                  d      d        Zy) z7Gradients for operators defined in control_flow_ops.py.    )dtypes)indexed_slices)ops)sparse_tensor)tensor)control_flow_ops)control_flow_util)math_ops)*c                 8   t        j                         }| j                         }|j                         }t        |t              r|j
                  j                  j                  |       }|!|d   t        j                  ||d   d       y|d   5t        |d   |d   gd      d   }||j
                  j                  | <   |dfS yt        |t              rt|d|j                  z
     }|N| j                  d   j                  t        j                   k(  r#t        ||j                     gd	z  d
      d   dfS yt        |d      d   dfS t#        |d   | j                  d         d   }t#        |d   | j                  d         d   }t        ||g      d   dfS )a  Gradients for a Switch op is calculated using a Merge op.

  If the switch is a loop switch, it will be visited twice. We create
  the merge on the first visit, and update the other input of the merge
  on the second visit. A next_iteration is also added on second visit.
  N   F)enforce_shape_invariant)NNr   b_switchname   cond_resource_grad	cond_grad)r   get_default_graph_get_control_flow_context
isinstanceWhileContext
grad_state
switch_mapgetr   _AddNextAndBackEdgemergeCondContextbranchinputsdtyper   resourceswitch)	opgradgraphop_ctxt	grad_ctxt
merge_grad	zero_grad
false_grad	true_grads	            W/home/dcms/DCMS/lib/python3.12/site-packages/tensorflow/python/ops/control_flow_grad.py_SwitchGradr.       s    


!%((*'--/)&%%0044R8J
 
a	,,ZaEJ	L 	a	
 $q'47+*=a@j,6i%%b)
 ';'Q'(I  
1		v	.'..!"Q&-ACCDFGKL 	LK(+T11Q1.q1JtAw		!-a0I*i()!,d22    Switch	RefSwitchMergec                    | j                   d   j                  }t        j                         }t	        j
                  |      }|j                         }t        |t              r t        j                  ||j                        S t        |t              r|j                  }|r|j                  r|j                  }|j                  j!                  |j"                        }	|	h|j$                  }|j'                          |j)                  |      }
|j+                          |j-                  |
|      }	|	|j                  |j"                  <   |	}t        j                  ||d      S t/        | j                         }t1        |      D cg c]%  }t3        j4                  | j6                  d   |      ' }}t1        |      D cg c]  }t        j                  |||         d     c}S c c}w c c}w )z:Gradients for a Merge op are calculated using a Switch op.r   r   r   r   )r    r$   r   r   r	   GetOutputContextr   r   r   r   _SwitchRefOrTensorpivotr   predr   history_mapr   r   grad_contextExitAddForwardAccumulatorEnterAddBackpropAccumulatedValuelenranger
   equaloutputs)r$   r%   _input_opr&   r'   r(   r7   r   	real_predhistory_pred
num_inputsiconds                 r-   
_MergeGradrI   \   s    YYq\__(



!%..x8'--/)&..tY__EE';'<<DY))
 ''j((,,TYY7i		++	!77= ::<N	,5
tyy)d..tTLL RYYJ6;J6GHHNN2::a=!,HDH z" 	++D$q':1=  Is   *G#=#G(RefMergec                     t        | ||      S N)rI   )r$   r%   rB   s      r-   _RefMergeGradrM      s    	Ba	  r/   r:   c                 ^   t        j                         }| j                         }|j                         }|j                  sy|j                  rt        d      t        |t        j                        r|j                  |j                         nt        |t        j                  t        j                  f      st        dt        |       d      |j                  |j                   j                         |j                  |j"                  j                         |j$                  }||j                  |j                         |j'                          t)        j*                  ||j                  d|j,                  d      }|j.                  j1                  |       |j3                          |S )z:Gradients for an exit op are calculated using an Enter op.Nz4Second-order gradient for while loops not supported.Type zO not supported, must be either`indexed_slices.IndexedSlices` or `SparseTensor`.Fb_exit)is_constantparallel_iterationsr   )r   r   r   	back_propr   	TypeErrorr   r   TensorAddNamer   r   IndexedSlicesr   SparseTensortypevaluesindicesdense_shaper<   r   _EnterrR   loop_entersappendr:   )r$   r%   r&   r'   r(   r\   results          r-   	_ExitGradra      s`    


!%((*'--/)			 
J
KKfmm$dii ~++]-G-GHJd4j\ *J J K Kdkk&&'dll''(""K(()//""
INN#77&
 v&..	-r/   RefExitNextIterationc                     |S )zA forward next_iteration is translated into a backprop identity.

  Note that the backprop next_iteration is added in switch grad.
   rB   r%   s     r-   _NextIterationGradrg      s	     
+r/   RefNextIterationc                     t        | |      S rL   )rg   rf   s     r-   _RefNextIterationGradrj      s    	At	$$r/   r<   c                    t        j                         }|j                         }||S |j                  s|S |j                  |S | j                  d      rtt        |t        j                        r|j                  | |      }|S t        |t        j                        r|j                  | |      }|S t        dt        |       d      t        |      }|j                   j#                  |       |j%                  |g       |S )zGradients for an Enter are calculated using an Exit op.

  For loop variables, grad is the gradient so just add an exit.
  For loop invariants, we need to add an accumulator loop.
  rQ   rO   z/ not supported,must be Tensor or Indexed Slices)r   r   r   rS   r   get_attrr   r   rU   AddBackpropAccumulatorr   rW   #AddBackpropIndexedSlicesAccumulatorrT   rY   exit
loop_exitsr_   
ExitResult)r$   r%   r&   r(   r`   s        r-   
_EnterGradrr      s     


!%--/)K			K!K[[$&//D9f 
- 
D.66	7<<RFf 
- d4j\ *9 9 : : $ZF'&"	-r/   RefEnterc                     t        | |      S rL   )rr   )r$   r%   s     r-   _RefEnterGradru      s    	B	r/   LoopCondc                      y)z0Stop backprop for the predicate of a while loop.Nre   )rB   s    r-   _LoopCondGradrx      s     
r/   N)__doc__tensorflow.python.frameworkr   r   r   r   r   tensorflow.python.opsr   r	   r
   &tensorflow.python.ops.control_flow_opsr.   RegisterGradientrI   rM   ra   rg   rj   rr   ru   rx   re   r/   r-   <module>r~      s    > . 6 + 5 . 2 3 * 553p    X { + !   [ !+ . g+ +^ j!! "! f% %P     Y 	 * o& ' ()% *% g   F j! " j! "r/   