@register_stabilize
@register_canonicalize
@local_optimizer([Solve])
def tag_solve_triangular(node):
 if (node.op == solve):
  if (node.op.A_structure == 'general'):
   (A, b) = node.inputs
   if (A.owner and isinstance(A.owner.op, type(cholesky))):
    if A.owner.op.lower:
     return [Solve('lower_triangular')(A, b)]
    else:
     return [Solve('upper_triangular')(A, b)]
   if (A.owner and isinstance(A.owner.op, DimShuffle) and (A.owner.op.new_order == (1, 0))):
    (A_T,) = A.owner.inputs
    if (A_T.owner and isinstance(A_T.owner.op, type(cholesky))):
     if A_T.owner.op.lower:
      return [Solve('upper_triangular')(A, b)]
     else:
      return [Solve('lower_triangular')(A, b)]

@ register_stabilize @ register_canonicalize @ local_optimizer ( [ Solve ] ) def tag_solve_triangular ( node ) : if ( node . op == solve ) : if ( node . op . A_structure == 'general' ) : ( A , b ) = node . inputs if ( A . owner and isinstance ( A . owner . op , type ( cholesky ) ) ) : if A . owner . op . lower : return [ Solve ( 'lower_triangular' ) ( A , b ) ] else : return [ Solve ( 'upper_triangular' ) ( A , b ) ] if ( A . owner and isinstance ( A . owner . op , DimShuffle ) and ( A . owner . op . new_order == ( 1 , 0 ) ) ) : ( A_T , ) = A . owner . inputs if ( A_T . owner and isinstance ( A_T . owner . op , type ( cholesky ) ) ) : if A_T . owner . op . lower : return [ Solve ( 'upper_triangular' ) ( A , b ) ] else : return [ Solve ( 'lower_triangular' ) ( A , b ) ]
if a general solve() is applied to the output of a cholesky op .