def _rotate_decomp_bare(U, P, A, iterations, backend):
  """Applies a small unitary rotation to a polar decomposition.

  A is an anti-Hermitian matrix that defines the unitary V = exp(-A). The polar
  decomposition (U, P) is then rotated to be (U @ V, V^dg @ P).

  This is done by expanding the Taylor series of exp(-A), assuming A is small.
  The order of the expansion and precision for the matmuls used at each order is
  encoded in `iterations`, which should be the output of
  `_rotate_decomp_preprocess`.

  This, the _bare version of the function, should be either jitted or pmapped
  with `static_argnums=(3, 4)`.

  Args:
    U, P: The polar decomposition to rotate.
    A: The anti-Hermitian matrix defining the rotation.
    iterations: The output of `_rotate_decomp_preprocess`, holding information
      about the orders in the expansion.
    backend: A `distla_core.linalg.backends.distributed_backend` or `serial_backend`

  Returns:
    U_rotated, P_rotated: The rotated polar decomposition.
  """
  # The cumulants will hold the current term in the expansion.
  U_cumulant = U
  P_cumulant = P
  input_dtype = U.dtype
  for (factor, sign, dtype, precision) in iterations:
    # dtypes at successive iterations only go down in precision, so we can
    # safely cast A and the cumulants.
    if A.dtype != dtype:
      A = A.astype(dtype)
    if U_cumulant.dtype != dtype:
      U_cumulant = U_cumulant.astype(dtype)
    if P_cumulant.dtype != dtype:
      P_cumulant = P_cumulant.astype(dtype)
    # The terms to add are U @ (-A)**order / order! and A**order @ P / order!
    U_cumulant = backend.matmul(U_cumulant, A, precision=precision)
    P_cumulant = backend.matmul(A, P_cumulant, precision=precision)
    U_term = dtype(sign / factor) * U_cumulant
    P_term = dtype(1 / factor) * P_cumulant
    U = U + U_term.astype(input_dtype)
    P = P + P_term.astype(input_dtype)
  return U, P