def det(r1, r2, r3):
  """Calculates the determinant of a 3x3 matrix with rows as args.

  Args:
    r1: First row
    r2: Second row
    r3: Third row

  Returns:
    Determinant of matrix [r1,r2,r3].  Functionally equivalent to
    jnp.linalg.det(jnp.array([r1,r2,r3])), but jits 10x faster for large batch.

  """
  return r1[0] * r2[1] * r3[2] + r1[1] * r2[2] * r3[0] + r1[2] * r2[0] * r3[
      1] - r1[2] * r2[1] * r3[0] - r1[0] * r2[2] * r3[1] - r1[1] * r2[0] * r3[2]