import torch


def cross_op_torch(r):
  """
  Return the cross operator as a matrix
  i.e. for input vector r \in \R^3
  output rX s.t. rX.dot(v) = np.cross(r, v)
  where rX \in \R^{3 X 3}
  """
  if len(r.shape) > 1:
      rX = torch.zeros(r.shape[0], 3, 3).to(r)
      rX[..., 0, 1] = -r[..., 2]
      rX[..., 0, 2] = r[..., 1]
      rX[..., 1, 2] = -r[..., 0]
      rX = rX - rX.transpose(-1, -2)
  else:
      rX = torch.zeros(3, 3).to(r)
      rX[0, 1] = -r[2]
      rX[0, 2] = r[1]
      rX[1, 2] = -r[0]
      rX = rX - rX.T
  return rX