def _compose_output_rep(lhs_rep, rhs_rep, lhs_contraction, rhs_contraction,
                        lhs_batch, rhs_batch):
  """Compose the output string representation.

  e.g., ij, jk, (((1,), (0,)), ((), ())) -> ik
        aij, ajk, (((2,), (1,)), ((0,), (0,))) -> aik

  Args:
    lhs_rep: A string representation for the left-hand side input array
    rhs_rep: A string representation for the right-hand side input array
    lhs_contraction: Sequence[int] (the contraction dimensions of lhs)
    rhs_contraction: Sequence[int] (the contraction dimensions of rhs)
    lhs_batch: Sequence[int] (the batch dimensions of lhs)
    rhs_batch: Sequence[int] (the batch dimensions of rhs)

  Returns:
    A string representation of the result array.
  """
  output_rep = []
  for dim in lhs_batch:
    output_rep.append(lhs_rep[dim])

  for i in _minus(range(len(lhs_rep)), lhs_batch + lhs_contraction):
    output_rep.append(lhs_rep[i])
  for i in _minus(range(len(rhs_rep)), rhs_batch + rhs_contraction):
    output_rep.append(rhs_rep[i])
  return "".join(output_rep)
Compose the output string representation.

  e.g., ij, jk, (((1,), (0,)), ((), ())) -> ik
        aij, ajk, (((2,), (1,)), ((0,), (0,))) -> aik

  Args:
    lhs_rep: A string representation for the left-hand side input array
    rhs_rep: A string representation for the right-hand side input array
    lhs_contraction: Sequence[int] (the contraction dimensions of lhs)
    rhs_contraction: Sequence[int] (the contraction dimensions of rhs)
    lhs_batch: Sequence[int] (the batch dimensions of lhs)
    rhs_batch: Sequence[int] (the batch dimensions of rhs)

  Returns:
    A string representation of the result array.
  
trax/tf_numpy/extensions/extensions.py