parameters: (matrices, powers)
return_statement: return array_ops.where(
        math_ops.equal(final_residual_power,
                       array_ops.zeros_like(
                           final_residual_power, dtype=dtypes.int32)),
        ident_tiled, math_ops.matmul(final_argument, final_accumulator))
docstring: Compute powers of matrices, e.g. A^3 = matmul(matmul(A, A), A).
  
    Uses exponentiation by squaring, with O(log(p)) matrix multiplications to
    compute A^p.
  
    Args:
      matrices: [batch size x N x N]
      powers: Which integer power to raise each matrix to [batch size]
    Returns:
      The matrices raised to their respective powers, same dimensions as the
      "matrices" argument.
docstring_summary: Compute powers of matrices, e.g. A^3 = matmul(matmul(A, A), A).
docstring_tokens: Compute powers of matrices e . g . A^3 = matmul ( matmul ( A A ) A ) .
function: def batch_matrix_pow(matrices, powers):
    """Compute powers of matrices, e.g. A^3 = matmul(matmul(A, A), A).
  
    Uses exponentiation by squaring, with O(log(p)) matrix multiplications to
    compute A^p.
  
    Args:
      matrices: [batch size x N x N]
      powers: Which integer power to raise each matrix to [batch size]
    Returns:
      The matrices raised to their respective powers, same dimensions as the
      "matrices" argument.
    """
  
    def terminate_when_all_zero(current_argument, residual_powers, accumulator):
      del current_argument, accumulator  # not used for condition
      do_exit = math_ops.reduce_any(
          math_ops.greater(residual_powers, array_ops.ones_like(residual_powers)))
      return do_exit
  
    def do_iteration(current_argument, residual_powers, accumulator):
      """Compute one step of iterative exponentiation by squaring.
  
      The recursive form is:
        power(A, p) = { power(matmul(A, A), p / 2) for even p
                      { matmul(A, power(matmul(A, A), (p - 1) / 2)) for odd p
        power(A, 0) = I
  
      The power(A, 0) = I case is handeled by starting with accumulator set to the
      identity matrix; matrices with zero residual powers are passed through
      unchanged.
  
      Args:
        current_argument: On this step, what is the first argument (A^2..^2) to
            the (unrolled) recursive function? [batch size x N x N]
        residual_powers: On this step, what is the second argument (residual p)?
            [batch_size]
        accumulator: Accumulates the exterior multiplications from the odd
            powers (initially the identity matrix). [batch_size x N x N]
      Returns:
        Updated versions of each argument for one step of the unrolled
        computation. Does not change parts of the batch which have a residual
        power of zero.
      """
      is_even = math_ops.equal(residual_powers % 2,
                               array_ops.zeros(
                                   array_ops.shape(residual_powers),
                                   dtype=dtypes.int32))
      new_accumulator = array_ops.where(is_even, accumulator,
                                        math_ops.matmul(accumulator,
                                                        current_argument))
      new_argument = math_ops.matmul(current_argument, current_argument)
      do_update = math_ops.greater(residual_powers, 1)
      new_residual_powers = residual_powers - residual_powers % 2
      new_residual_powers //= 2
      # Stop updating if we've reached our base case; some batch elements may
      # finish sooner than others
      accumulator = array_ops.where(do_update, new_accumulator, accumulator)
      current_argument = array_ops.where(do_update, new_argument,
                                         current_argument)
      residual_powers = array_ops.where(do_update, new_residual_powers,
                                        residual_powers)
      return (current_argument, residual_powers, accumulator)
  
    matrices = ops.convert_to_tensor(matrices)
    powers = math_ops.cast(powers, dtype=dtypes.int32)
    ident = array_ops.expand_dims(
        array_ops.diag(
            array_ops.ones([array_ops.shape(matrices)[1]], dtype=matrices.dtype)),
        0)
    ident_tiled = array_ops.tile(ident, [array_ops.shape(matrices)[0], 1, 1])
    (final_argument,
     final_residual_power, final_accumulator) = control_flow_ops.while_loop(
         terminate_when_all_zero, do_iteration, [matrices, powers, ident_tiled])
    return array_ops.where(
        math_ops.equal(final_residual_power,
                       array_ops.zeros_like(
                           final_residual_power, dtype=dtypes.int32)),
        ident_tiled, math_ops.matmul(final_argument, final_accumulator))
function_tokens: def batch_matrix_pow ( matrices , powers ) : def terminate_when_all_zero ( current_argument , residual_powers , accumulator ) : del current_argument , accumulator # not used for condition do_exit = math_ops . reduce_any ( math_ops . greater ( residual_powers , array_ops . ones_like ( residual_powers ) ) ) return do_exit def do_iteration ( current_argument , residual_powers , accumulator ) : """Compute one step of iterative exponentiation by squaring.
  
      The recursive form is:
        power(A, p) = { power(matmul(A, A), p / 2) for even p
                      { matmul(A, power(matmul(A, A), (p - 1) / 2)) for odd p
        power(A, 0) = I
  
      The power(A, 0) = I case is handeled by starting with accumulator set to the
      identity matrix; matrices with zero residual powers are passed through
      unchanged.
  
      Args:
        current_argument: On this step, what is the first argument (A^2..^2) to
            the (unrolled) recursive function? [batch size x N x N]
        residual_powers: On this step, what is the second argument (residual p)?
            [batch_size]
        accumulator: Accumulates the exterior multiplications from the odd
            powers (initially the identity matrix). [batch_size x N x N]
      Returns:
        Updated versions of each argument for one step of the unrolled
        computation. Does not change parts of the batch which have a residual
        power of zero.
      """ is_even = math_ops . equal ( residual_powers % 2 , array_ops . zeros ( array_ops . shape ( residual_powers ) , dtype = dtypes . int32 ) ) new_accumulator = array_ops . where ( is_even , accumulator , math_ops . matmul ( accumulator , current_argument ) ) new_argument = math_ops . matmul ( current_argument , current_argument ) do_update = math_ops . greater ( residual_powers , 1 ) new_residual_powers = residual_powers - residual_powers % 2 new_residual_powers //= 2 # Stop updating if we've reached our base case; some batch elements may # finish sooner than others accumulator = array_ops . where ( do_update , new_accumulator , accumulator ) current_argument = array_ops . where ( do_update , new_argument , current_argument ) residual_powers = array_ops . where ( do_update , new_residual_powers , residual_powers ) return ( current_argument , residual_powers , accumulator ) matrices = ops . convert_to_tensor ( matrices ) powers = math_ops . cast ( powers , dtype = dtypes . int32 ) ident = array_ops . expand_dims ( array_ops . diag ( array_ops . ones ( [ array_ops . shape ( matrices ) [ 1 ] ] , dtype = matrices . dtype ) ) , 0 ) ident_tiled = array_ops . tile ( ident , [ array_ops . shape ( matrices ) [ 0 ] , 1 , 1 ] ) ( final_argument , final_residual_power , final_accumulator ) = control_flow_ops . while_loop ( terminate_when_all_zero , do_iteration , [ matrices , powers , ident_tiled ] ) return array_ops . where ( math_ops . equal ( final_residual_power , array_ops . zeros_like ( final_residual_power , dtype = dtypes . int32 ) ) , ident_tiled , math_ops . matmul ( final_argument , final_accumulator ) )