import torch


def homogenize(m):
  """Adds homogeneous coordinates to a [..., N,N] matrix, returning [..., N+1, N+1]."""
  assert m.shape[-1] == m.shape[-2]  # Must be square
  n = m.shape[-1]
  eye_n_plus_1 = torch.eye(n+1).cuda().expand(list(m.shape[:-2]) + [-1, -1])
  extra_col = eye_n_plus_1[..., :-1, -1:]
  extra_row = eye_n_plus_1[..., -1:, :]
  including_col = torch.cat([m, extra_col], dim=-1)
  return torch.cat([including_col, extra_row], dim=-2)