jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit(self, source_data: torch.Tensor, target_data: torch.Tensor) -> Mapping[str, Any]:
        translation_matrix = torch.linalg.lstsq(source_data, target_data).solution
        translation_matrix = torch.as_tensor(translation_matrix, dtype=source_data.dtype, device=source_data.device)
        U, _, Vt = torch.linalg.svd(translation_matrix)
        translation_matrix = U @ Vt.T
        self.register_buffer("translation_matrix", translation_matrix)
        return {}
//  [python]
def forward(ctx, A):
		b = int(A.shape[0])
		n = int(A.shape[1])
		u = []; s = []
		for _ in A.view(b,n,-1).detach():
			u_, s_, v_ = torch.svd(_)
			u.append(u_.unsqueeze(0))
			s.append(s_.unsqueeze(0))
		u = torch.cat(u, 0)
		s = torch.cat(s, 0)
		i = torch.argmax(s, 1)
		s = torch.gather(s, 1, i.view(-1,1)).squeeze(-1)
		u = torch.gather(u, 2, \
			i.view(-1,1,1).expand(-1,n,1))
		u = torch.where(u[:,-1:,:] \
			.expand(-1,n,-1) < 0, -u, u) \
			.squeeze(-1)
		ctx.save_for_backward(u, s, A)
		return u, s
//  [python]
def forward(ctx, A):
		b = int(A.shape[0])
		n = int(A.shape[1])
		u = []; s = []
		for _ in A.view(b,n,-1).detach():
			u_, s_, v_ = torch.linalg.svd(_)
			u.append(u_.unsqueeze(0))
			s.append(s_.unsqueeze(0))
		u = torch.cat(u, 0)
		s = torch.cat(s, 0)
		i = torch.argmax(s, 1)
		s = torch.gather(s, 1, i.view(-1,1)).squeeze(-1)
		u = torch.gather(u, 2, \
			i.view(-1,1,1).expand(-1,n,1))
		u = torch.where(u[:,-1:,:] \
			.expand(-1,n,-1) < 0, -u, u) \
			.squeeze(-1)
		ctx.save_for_backward(u, s, A)
		return u, s
//  [python]
def _make_params(self, random_init = True):
		w = getattr(self.module, self.name)

		height = w.data.shape[0]
		width = w.view(height, -1).data.shape[1]
		if self.power_iterations > 0:
			if random_init:
				u = w.data.new(height).normal_(0, 1)
				v = w.data.new(width).normal_(0, 1)
				u = Normalize.apply(u)
				v = Normalize.apply(v)
			else:
				u, _, v = torch.svd(w.data.view(height,-1))
				i = torch.argmax(_, 0)
				u = u[:, i]
				v = v[:, i]
			u = torch.nn.Parameter(u, requires_grad = False)
			v = torch.nn.Parameter(v, requires_grad = False)
			self.module.register_parameter(self.name + '_u', u)
			self.module.register_parameter(self.name + '_v', v)
		w_bar = torch.nn.Parameter(w.data)
		del self.module._parameters[self.name]
		self.module.register_parameter(self.name + '_bar', w_bar)
//  [python]
def _make_params(self, random_init = True):
		w = getattr(self.module, self.name)

		height = w.data.shape[0]
		width = w.view(height, -1).data.shape[1]
		if self.power_iterations > 0:
			if random_init:
				u = w.data.new(height).normal_(0, 1)
				v = w.data.new(width).normal_(0, 1)
				u = Normalize.apply(u)
				v = Normalize.apply(v)
			else:
				u, s, v = torch.linalg.svd(w.data.view(height,-1))
				i = torch.argmax(s, 0)
				u = u[:, i]
				v = v[:, i]
			u = torch.nn.Parameter(u, requires_grad = False)
			v = torch.nn.Parameter(v, requires_grad = False)
			self.module.register_parameter(self.name + '_u', u)
			self.module.register_parameter(self.name + '_v', v)
		w_bar = torch.nn.Parameter(w.data)
		del self.module._parameters[self.name]
		self.module.register_parameter(self.name + '_bar', w_bar)
//  [python]
def estimate_pose(self, pt0, pt1):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        cent0 = torch.sum(pt0 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()
        cent1 = torch.sum(pt1 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()

        diag_mat = torch.diag(self.pconf).unsqueeze(0)
        x = (pt0 - cent0).transpose(2, 1).contiguous()
        y = pt1 - cent1

        pred_t = cent1 - cent0

        cov = torch.bmm(torch.bmm(x, diag_mat), y).contiguous().squeeze(0)

        u, _, v = torch.svd(cov)

        u = u.transpose(1, 0).contiguous()
        d = torch.det(torch.mm(v, u)).contiguous().view(1, 1, 1).contiguous()
        u = u.transpose(1, 0).contiguous().unsqueeze(0)

        ud = torch.cat((u[:, :, :-1], u[:, :, -1:] * d), dim=2)
        v = v.transpose(1, 0).contiguous().unsqueeze(0)

        pred_r = torch.bmm(ud, v).transpose(2, 1).contiguous()
        return pred_r, pred_t[:, 0, :].view(1, 3)
//  [python]
def estimate_pose(self, pt0, pt1):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        cent0 = torch.sum(pt0 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()
        cent1 = torch.sum(pt1 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()

        diag_mat = torch.diag(self.pconf).unsqueeze(0)
        x = (pt0 - cent0).transpose(2, 1).contiguous()
        y = pt1 - cent1

        pred_t = cent1 - cent0

        cov = torch.bmm(torch.bmm(x, diag_mat), y).contiguous().squeeze(0)

        u, s, v = torch.linalg.svd(cov)

        u = u.transpose(1, 0).contiguous()
        d = torch.det(torch.mm(v, u)).contiguous().view(1, 1, 1).contiguous()
        u = u.transpose(1, 0).contiguous().unsqueeze(0)

        ud = torch.cat((u[:, :, :-1], u[:, :, -1:] * d), dim=2)
        v = v.transpose(1, 0).contiguous().unsqueeze(0)

        pred_r = torch.bmm(ud, v).transpose(2, 1).contiguous()
        return pred_r, pred_t[:, 0, :].view(1, 3)
//  [python]
def estimate_rotation(self, pt0, pt1, sym_or_not):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        cent0 = torch.sum(pt0 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()
        cent1 = torch.sum(pt1 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()

        diag_mat = torch.diag(self.pconf).unsqueeze(0)
        x = (pt0 - cent0).transpose(2, 1).contiguous()
        y = pt1 - cent1

        pred_t = cent1 - cent0

        cov = torch.bmm(torch.bmm(x, diag_mat), y).contiguous().squeeze(0)

        u, _, v = torch.svd(cov)

        u = u.transpose(1, 0).contiguous()
        d = torch.det(torch.mm(v, u)).contiguous().view(1, 1, 1).contiguous()
        u = u.transpose(1, 0).contiguous().unsqueeze(0)

        ud = torch.cat((u[:, :, :-1], u[:, :, -1:] * d), dim=2)
        v = v.transpose(1, 0).contiguous().unsqueeze(0)

        pred_r = torch.bmm(ud, v).transpose(2, 1).contiguous()

        if sym_or_not:
            pred_r = torch.bmm(pred_r, self.sym_axis).contiguous().view(-1).contiguous()

        return pred_r
//  [python]
def estimate_rotation(self, pt0, pt1, sym_or_not):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        cent0 = torch.sum(pt0 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()
        cent1 = torch.sum(pt1 * pconf2, dim=1).repeat(1, self.num_key, 1).contiguous()

        diag_mat = torch.diag(self.pconf).unsqueeze(0)
        x = (pt0 - cent0).transpose(2, 1).contiguous()
        y = pt1 - cent1

        pred_t = cent1 - cent0

        cov = torch.bmm(torch.bmm(x, diag_mat), y).contiguous().squeeze(0)

        u, s, v = torch.linalg.svd(cov)

        u = u.transpose(1, 0).contiguous()
        d = torch.det(torch.mm(v, u)).contiguous().view(1, 1, 1).contiguous()
        u = u.transpose(1, 0).contiguous().unsqueeze(0)

        ud = torch.cat((u[:, :, :-1], u[:, :, -1:] * d), dim=2)
        v = v.transpose(1, 0).contiguous().unsqueeze(0)

        pred_r = torch.bmm(ud, v).transpose(2, 1).contiguous()

        if sym_or_not:
            pred_r = torch.bmm(pred_r, self.sym_axis).contiguous().view(-1).contiguous()

        return pred_r
//  [python]
def fit(self, source_data: torch.Tensor, target_data: torch.Tensor) -> Mapping[str, Any]:
        self.dim_matcher.fit(source_data, target_data)
        dim_matcher_out = self.dim_matcher(source_data, target_data)
        source_data, target_data = dim_matcher_out["source"], dim_matcher_out["target"]

        assert source_data.size(1) == target_data.size(
            1
        ), f"Dimension mismatch between {source_data.size(1)} and {target_data.size(1)}. Forgot some padding/truncation transforms?"

        #  Compute the translation vector that aligns A to B using SVD.
        u, sigma, vt = torch.svd((target_data.T @ source_data).T)
        translation_matrix = u @ vt.T

        translation_matrix = torch.as_tensor(translation_matrix, dtype=source_data.dtype, device=source_data.device)
        self.register_buffer("translation_matrix", translation_matrix)

        sigma_rank = (~sigma.isclose(torch.zeros_like(sigma), atol=1e-1).bool()).sum().item()

        return {"sigma_rank": sigma_rank}
//  [python]
def fit(self, source_data: torch.Tensor, target_data: torch.Tensor) -> Mapping[str, Any]:
        self.dim_matcher.fit(source_data, target_data)
        dim_matcher_out = self.dim_matcher(source_data, target_data)
        source_data, target_data = dim_matcher_out["source"], dim_matcher_out["target"]

        assert source_data.size(1) == target_data.size(
            1
        ), f"Dimension mismatch between {source_data.size(1)} and {target_data.size(1)}. Forgot some padding/truncation transforms?"

        #  Compute the translation vector that aligns A to B using SVD.
        u, s, vt = torch.linalg.svd((target_data.T @ source_data).T)
        translation_matrix = u @ vt.T

        translation_matrix = torch.as_tensor(translation_matrix, dtype=source_data.dtype, device=source_data.device)
        self.register_buffer("translation_matrix", translation_matrix)

        sigma_rank = (~s.isclose(torch.zeros_like(s), atol=1e-1).bool()).sum().item()

        return {"sigma_rank": sigma_rank}
//  [python]
def forward(self, A):
        U, S, V = torch.svd(A) #NOTE: torch.svd may be depreciated later, requiring the switch to torch.linalg.svd
        self.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def forward(self, A):
        U, S, V = torch.linalg.svd(A)
        self.save_for_backward(U, S, V)
        return U, S, V
//  [python]
print(torch.cross(v2, v1)) # negative of z unit vector (v1 x v2 == -v2 x v1)
print(m1)
m3 = torch.matmul(m1, m2)
print(m3)                  # 3 times m1
print(torch.svd(m3))
//  [python]
print(torch.cross(v2, v1)) # negative of z unit vector (v1 x v2 == -v2 x v1)
print(m1)
m3 = torch.matmul(m1, m2)
print(m3)                  # 3 times m1
print(torch.linalg.svd(m3))
//  [python]
def elastic_crit(ret, batch):
    """
    resd_jacobian: n_batch, n_point, 3, 3
    """
    jac = ret['resd_jacobian']
    U, S, V = torch.svd(jac, compute_uv=True)
    log_svals = torch.log(torch.clamp(S, min=1e-6))
    elastic_loss = torch.sum(log_svals**2, dim=2).mean()
    return elastic_loss
//  [python]
def elastic_crit(ret, batch):
    """
    resd_jacobian: n_batch, n_point, 3, 3
    """
    jac = ret['resd_jacobian']
    U, S, V = torch.linalg.svd(jac)
    log_svals = torch.log(torch.clamp(S, min=1e-6))
    elastic_loss = torch.sum(log_svals**2, dim=2).mean()
    return elastic_loss
//  [python]
def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
        assert img_dim % ratio == 0
        self.img_dim = img_dim
        self.channels = channels
        self.y_dim = img_dim // ratio
        self.ratio = ratio
        H = torch.Tensor([[1 / ratio**2] * ratio**2]).to(device)
        self.U_small, self.singulars_small, self.V_small = torch.svd(H, some=False)
        self.Vt_small = self.V_small.transpose(0, 1)
//  [python]
def __init__(self, channels, img_dim, ratio, device): #ratio = 2 or 4
        assert img_dim % ratio == 0
        self.img_dim = img_dim
        self.channels = channels
        self.y_dim = img_dim // ratio
        self.ratio = ratio
        H = torch.Tensor([[1 / ratio**2] * ratio**2]).to(device)
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(H)
        self.Vt_small = self.V_small.transpose(0, 1)
//  [python]
def orthogonal_init(tensor, gain=1):
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    u, s, v = torch.svd(flattened, some=True)
    if rows < cols:
        u.t_()
    q = u if tuple(u.shape) == (rows, cols) else v
    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_init(tensor, gain=1):
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    u, s, v = torch.linalg.svd(flattened)
    if rows < cols:
        u.t_()
    q = u if tuple(u.shape) == (rows, cols) else v
    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def low_rank_approximation(self, matrix: Tensor) -> Tensor:
        """
        Helper function to perform low-rank approximation on a 2D matrix.

        Args:
            matrix (Tensor): The 2D matrix to be approximated.

        Returns:
            torch.Tensor: The approximated 2D matrix with reduced rank.
        """
        U, S, V = torch.svd(matrix)
        max_rank = min(matrix.size())
        approx_rank = int(self.rank_fraction * max_rank)
        U_r = U[:, :approx_rank]
        S_r = S[:approx_rank]
        V_r = V[:, :approx_rank]
        W_approx = torch.mm(U_r, torch.mm(torch.diag(S_r), V_r.t()))
        return W_approx
//  [python]
def low_rank_approximation(self, matrix: Tensor) -> Tensor:
        """
        Helper function to perform low-rank approximation on a 2D matrix.

        Args:
            matrix (Tensor): The 2D matrix to be approximated.

        Returns:
            torch.Tensor: The approximated 2D matrix with reduced rank.
        """
        U, S, V = torch.linalg.svd(matrix)
        max_rank = min(matrix.size())
        approx_rank = int(self.rank_fraction * max_rank)
        U_r = U[:, :approx_rank]
        S_r = S[:approx_rank]
        V_r = V[:, :approx_rank]
        W_approx = torch.mm(U_r, torch.mm(torch.diag(S_r), V_r.t()))
        return W_approx
//  [python]
def compute_rigid_transform(a, b, weights):
    """Compute rigid transforms between two point sets

    Args:
        a (torch.Tensor): (B, M, 3) points
        b (torch.Tensor): (B, N, 3) points
        weights (torch.Tensor): (B, M)

    Returns:
        Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
    """

    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
    centroid_a = torch.sum(a * weights_normalized, dim=1)
    centroid_b = torch.sum(b * weights_normalized, dim=1)
    a_centered = a - centroid_a[:, None, :]
    b_centered = b - centroid_b[:, None, :]
    cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)

    # Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
    # and choose based on determinant to avoid flips
    u, s, v = torch.svd(cov, some=False, compute_uv=True)
    rot_mat_pos = v @ u.transpose(-1, -2)
    v_neg = v.clone()
    v_neg[:, :, 2] *= -1
    rot_mat_neg = v_neg @ u.transpose(-1, -2)
    rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
    assert torch.all(torch.det(rot_mat) > 0)

    # Compute translation (uncenter centroid)
    translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]

    transform = torch.cat((rot_mat, translation), dim=2)
    tsfm = torch.eye(4)
    tsfm[:3] = transform
    return tsfm.numpy()
//  [python]
def compute_rigid_transform(a, b, weights):
    """Compute rigid transforms between two point sets

    Args:
        a (torch.Tensor): (B, M, 3) points
        b (torch.Tensor): (B, N, 3) points
        weights (torch.Tensor): (B, M)

    Returns:
        Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
    """

    weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
    centroid_a = torch.sum(a * weights_normalized, dim=1)
    centroid_b = torch.sum(b * weights_normalized, dim=1)
    a_centered = a - centroid_a[:, None, :]
    b_centered = b - centroid_b[:, None, :]
    cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)

    # Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
    # and choose based on determinant to avoid flips
    u, s, v = torch.linalg.svd(cov)
    rot_mat_pos = v @ u.transpose(-1, -2)
    v_neg = v.clone()
    v_neg[:, :, 2] *= -1
    rot_mat_neg = v_neg @ u.transpose(-1, -2)
    rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
    assert torch.all(torch.det(rot_mat) > 0)

    # Compute translation (uncenter centroid)
    translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]

    transform = torch.cat((rot_mat, translation), dim=2)
    tsfm = torch.eye(4)
    tsfm[:3] = transform
    return tsfm.numpy()
//  [python]
def solve_procrustes_svd(pose, sources, targets, weights=None, **kwargs):
    """
    solve point-to-point procrustes using the Kabsch algorithm
    :param pose: initial pose (alpha, beta, gamma, tx, ty, tz), shape 6
    :param sources: source points, shape Kx3
    :param targets: target points, shape Kx3
    :param weights: correspondence weights, shape Kx1
    :param kwargs: unused keyword args
    :return: best transformation between source and target, shape 6
    """
    source_center = weighted_centroid(sources, weights)
    target_center = weighted_centroid(targets, weights)
    sources = sources - source_center
    targets = targets - target_center
    if weights is not None:
        svd_w = torch.matmul(torch.diag(weights.view(-1)), targets)
        svd_w = torch.matmul(sources.t(), svd_w)
    else:
        svd_w = torch.matmul(sources.t(), targets)  # 3x3
    check_conditioning(svd_w)
    svd_u, _, svd_v = torch.svd(svd_w)  # 3x3, 1x3, 3x3
    rotation = torch.matmul(svd_v, svd_u.t())  # 3x3
    if torch.det(rotation) < 0:  # reflection check
        reflection_fix = torch.ones_like(svd_v, dtype=svd_v.dtype, device=svd_v.device)
        reflection_fix[:, -1] = -1
        svd_v = svd_v * reflection_fix
        rotation = torch.matmul(svd_v, svd_u.t())
    translation = target_center.t() - torch.matmul(rotation, source_center.t())  # 1x3
    pose_change = torch.zeros_like(pose)  # 4x4
    pose_change[:3, :3] = rotation
    pose_change[:3, 3] = translation.view(-1)
    pose_change[3, 3] = 1
    return torch.matmul(pose_change, pose)
//  [python]
def solve_procrustes_svd(pose, sources, targets, weights=None, **kwargs):
    """
    solve point-to-point procrustes using the Kabsch algorithm
    :param pose: initial pose (alpha, beta, gamma, tx, ty, tz), shape 6
    :param sources: source points, shape Kx3
    :param targets: target points, shape Kx3
    :param weights: correspondence weights, shape Kx1
    :param kwargs: unused keyword args
    :return: best transformation between source and target, shape 6
    """
    source_center = weighted_centroid(sources, weights)
    target_center = weighted_centroid(targets, weights)
    sources = sources - source_center
    targets = targets - target_center
    if weights is not None:
        svd_w = torch.matmul(torch.diag(weights.view(-1)), targets)
        svd_w = torch.matmul(sources.t(), svd_w)
    else:
        svd_w = torch.matmul(sources.t(), targets)  # 3x3
    check_conditioning(svd_w)
    svd_u, s, svd_v = torch.linalg.svd(svd_w)  # 3x3, 1x3, 3x3
    rotation = torch.matmul(svd_v, svd_u.t())  # 3x3
    if torch.det(rotation) < 0:  # reflection check
        reflection_fix = torch.ones_like(svd_v, dtype=svd_v.dtype, device=svd_v.device)
        reflection_fix[:, -1] = -1
        svd_v = svd_v * reflection_fix
        rotation = torch.matmul(svd_v, svd_u.t())
    translation = target_center.t() - torch.matmul(rotation, source_center.t())  # 1x3
    pose_change = torch.zeros_like(pose)  # 4x4
    pose_change[:3, :3] = rotation
    pose_change[:3, 3] = translation.view(-1)
    pose_change[3, 3] = 1
    return torch.matmul(pose_change, pose)
//  [python]
def test_SVDARNOLDI_rank_deficient():
    m= 50
    k=15
    for r in [25,35,40,45]:
        M= torch.rand((m,m),dtype=torch.float64)
        U, S0, V= torch.svd(M)
        S0[-r:]=0
        M= U@torch.diag(S0)@V.t()

        U, S, V= SVDARNOLDI.apply(M, k)
        assert( torch.norm(M-U@torch.diag(S)@V.t())-torch.sqrt(torch.sum(S0[k:]**2)) 
            < S0[0]*(m**2)*1e-14 )
//  [python]
def test_SVDARNOLDI_rank_deficient():
    m= 50
    k=15
    for r in [25,35,40,45]:
        M= torch.rand((m,m),dtype=torch.float64)
        U, S0, V= torch.linalg.svd(M)
        S0[-r:]=0
        M= U@torch.diag(S0)@V.t()

        U, S, V= SVDARNOLDI.apply(M, k)
        assert( torch.norm(M-U@torch.diag(S)@V.t())-torch.sqrt(torch.sum(S0[k:]**2)) 
            < S0[0]*(m**2)*1e-14 )
//  [python]
def test_SVDARNOLDI_random():
    m= 50
    k= 10
    M= torch.rand(m, m, dtype=torch.float64)

    U0, S0, V0= torch.svd(M)

    U,S,V= SVDARNOLDI.apply(M,k)
    # |M|=\sqrt{Tr(MM^t)}=\sqrt{Tr(D^2)} => 
    # |M-US_kV^t|=\sqrt{Tr(D^2)-Tr(S^2)}=\sqrt{\sum_i>k D^2_i}
    assert( torch.norm(M-U@torch.diag(S)@V.t())-torch.sqrt(torch.sum(S0[k:]**2)) 
        < S0[0]*(m**2)*1e-14 )
//  [python]
def test_SVDARNOLDI_random():
    m= 50
    k= 10
    M= torch.rand(m, m, dtype=torch.float64)

    U0, S0, V0= torch.linalg.svd(M)

    U,S,V= SVDARNOLDI.apply(M,k)
    # |M|=\sqrt{Tr(MM^t)}=\sqrt{Tr(D^2)} => 
    # |M-US_kV^t|=\sqrt{Tr(D^2)-Tr(S^2)}=\sqrt{\sum_i>k D^2_i}
    assert( torch.norm(M-U@torch.diag(S)@V.t())-torch.sqrt(torch.sum(S0[k:]**2)) 
        < S0[0]*(m**2)*1e-14 )
//  [python]
def _torch_svd_cast(input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Helper function to make torch.svd work with other than fp32/64.

    The function torch.svd is only implemented for fp32/64 which makes
    impossible to be used by fp16 or others. What this function does, is cast
    input data type to fp32, apply torch.svd, and cast back to the input dtype.

    NOTE: in torch 1.8.1 this function is recommended to use as torch.linalg.svd
    """
    if not isinstance(input, torch.Tensor):
        raise AssertionError(f"Input must be torch.Tensor. Got: {type(input)}.")
    dtype: torch.dtype = input.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32

    out1, out2, out3 = torch.svd(input.to(dtype))

    return (out1.to(input.dtype), out2.to(input.dtype), out3.to(input.dtype))
//  [python]
def _torch_svd_cast(input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Helper function to make torch.svd work with other than fp32/64.

    The function torch.svd is only implemented for fp32/64 which makes
    impossible to be used by fp16 or others. What this function does, is cast
    input data type to fp32, apply torch.svd, and cast back to the input dtype.

    NOTE: in torch 1.8.1 this function is recommended to use as torch.linalg.svd
    """
    if not isinstance(input, torch.Tensor):
        raise AssertionError(f"Input must be torch.Tensor. Got: {type(input)}.")
    dtype: torch.dtype = input.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32

    out1, out2, out3 = torch.linalg.svd(input.to(dtype))

    return (out1.to(input.dtype), out2.to(input.dtype), out3.to(input.dtype))
//  [python]
def _r2(X, Y, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    UX, sX, _ = torch.svd(X)
    UY, sY, _ = torch.svd(Y)
    return _r1(UX @ torch.diag(sX), UY @ torch.diag(sY))
//  [python]
def _r2(X, Y, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    UX, sX, _ = torch.linalg.svd(X)
    UY, sY, _ = torch.linalg.svd(Y)
    return _r1(UX @ torch.diag(sX), UY @ torch.diag(sY))
//  [python]
def __init__(self, kernel, channels, img_dim, device, stride = 1):
        self.img_dim = img_dim
        self.channels = channels
        self.ratio = stride
        small_dim = img_dim // stride
        self.small_dim = small_dim
        #build 1D conv matrix
        H_small = torch.zeros(small_dim, img_dim, device=device)
        for i in range(stride//2, img_dim + stride//2, stride):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                j_effective = j
                #reflective padding
                if j_effective < 0: j_effective = -j_effective-1
                if j_effective >= img_dim: j_effective = (img_dim - 1) - (j_effective - img_dim)
                #matrix building
                H_small[i // stride, j_effective] += kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.svd(H_small, some=False)
        ZERO = 3e-2
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small.reshape(small_dim, 1), self.singulars_small.reshape(1, small_dim)).reshape(small_dim**2)
        #permutation for matching the singular values. See P_1 in Appendix D.5.
        self._perm = torch.Tensor([self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim)] + \
                                  [self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim, self.img_dim)]).to(device).long()
//  [python]
def __init__(self, kernel, channels, img_dim, device, stride = 1):
        self.img_dim = img_dim
        self.channels = channels
        self.ratio = stride
        small_dim = img_dim // stride
        self.small_dim = small_dim
        #build 1D conv matrix
        H_small = torch.zeros(small_dim, img_dim, device=device)
        for i in range(stride//2, img_dim + stride//2, stride):
            for j in range(i - kernel.shape[0]//2, i + kernel.shape[0]//2):
                j_effective = j
                #reflective padding
                if j_effective < 0: j_effective = -j_effective-1
                if j_effective >= img_dim: j_effective = (img_dim - 1) - (j_effective - img_dim)
                #matrix building
                H_small[i // stride, j_effective] += kernel[j - i + kernel.shape[0]//2]
        #get the svd of the 1D conv
        self.U_small, self.singulars_small, self.V_small = torch.linalg.svd(H_small)
        ZERO = 3e-2
        self.singulars_small[self.singulars_small < ZERO] = 0
        #calculate the singular values of the big matrix
        self._singulars = torch.matmul(self.singulars_small.reshape(small_dim, 1), self.singulars_small.reshape(1, small_dim)).reshape(small_dim**2)
        #permutation for matching the singular values. See P_1 in Appendix D.5.
        self._perm = torch.Tensor([self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim)] + \
                                  [self.img_dim * i + j for i in range(self.small_dim) for j in range(self.small_dim, self.img_dim)]).to(device).long()
//  [python]
def test_svd_lowrank(self, device, dtype):
        from torch.testing._internal.common_utils import random_lowrank_matrix, random_sparse_matrix

        def run_subtest(actual_rank, matrix_size, batches, device, svd_lowrank, **options):
            density = options.pop('density', 1)
            if isinstance(matrix_size, int):
                rows = columns = matrix_size
            else:
                rows, columns = matrix_size
            if density == 1:
                a_input = random_lowrank_matrix(actual_rank, rows, columns, *batches, device=device, dtype=dtype)
                a = a_input
            else:
                assert batches == ()
                a_input = random_sparse_matrix(rows, columns, density, device=device, dtype=dtype)
                a = a_input.to_dense()

            q = min(*size)
            u, s, v = svd_lowrank(a_input, q=q, **options)

            # check if u, s, v is a SVD
            u, s, v = u[..., :q], s[..., :q], v[..., :q]
            A = u.matmul(s.diag_embed()).matmul(v.mT)
            self.assertEqual(A, a, rtol=1e-7, atol=2e-7)

            # check if svd_lowrank produces same singular values as torch.svd
            U, S, V = torch.svd(a)
            self.assertEqual(s.shape, S.shape)
            self.assertEqual(u.shape, U.shape)
            self.assertEqual(v.shape, V.shape)
            self.assertEqual(s, S)

            if density == 1:
                # actual_rank is known only for dense inputs
                #
                # check if pairs (u, U) and (v, V) span the same
                # subspaces, respectively
                u, s, v = u[..., :actual_rank], s[..., :actual_rank], v[..., :actual_rank]
                U, S, V = U[..., :actual_rank], S[..., :actual_rank], V[..., :actual_rank]
                self.assertEqual(u.mT.matmul(U).det().abs(), torch.ones(batches, device=device, dtype=dtype))
                self.assertEqual(v.mT.matmul(V).det().abs(), torch.ones(batches, device=device, dtype=dtype))

        all_batches = [(), (1,), (3,), (2, 3)]
        for actual_rank, size, all_batches in [
                (2, (17, 4), all_batches),
                (4, (17, 4), all_batches),
                (4, (17, 17), all_batches),
                (10, (100, 40), all_batches),
                (7, (1000, 1000), [()]),
        ]:
            # dense input
            for batches in all_batches:
                run_subtest(actual_rank, size, batches, device, torch.svd_lowrank)
                if size != size[::-1]:
                    run_subtest(actual_rank, size[::-1], batches, device, torch.svd_lowrank)

        # sparse input
        for size in [(17, 4), (4, 17), (17, 17), (100, 40), (40, 100), (1000, 1000)]:
            for density in [0.005, 0.1]:
                run_subtest(None, size, (), device, torch.svd_lowrank, density=density)

        # jitting support
        jitted = torch.jit.script(torch.svd_lowrank)
        actual_rank, size, batches = 2, (17, 4), ()
        run_subtest(actual_rank, size, batches, device, jitted)
//  [python]
def test_svd_lowrank(self, device, dtype):
        from torch.testing._internal.common_utils import random_lowrank_matrix, random_sparse_matrix

        def run_subtest(actual_rank, matrix_size, batches, device, svd_lowrank, **options):
            density = options.pop('density', 1)
            if isinstance(matrix_size, int):
                rows = columns = matrix_size
            else:
                rows, columns = matrix_size
            if density == 1:
                a_input = random_lowrank_matrix(actual_rank, rows, columns, *batches, device=device, dtype=dtype)
                a = a_input
            else:
                assert batches == ()
                a_input = random_sparse_matrix(rows, columns, density, device=device, dtype=dtype)
                a = a_input.to_dense()

            q = min(*size)
            u, s, v = svd_lowrank(a_input, q=q, **options)

            # check if u, s, v is a SVD
            u, s, v = u[..., :q], s[..., :q], v[..., :q]
            A = u.matmul(s.diag_embed()).matmul(v.mT)
            self.assertEqual(A, a, rtol=1e-7, atol=2e-7)

            # check if svd_lowrank produces same singular values as torch.linalg.svd
            U, S, V = torch.linalg.svd(a)
            self.assertEqual(s.shape, S.shape)
            self.assertEqual(u.shape, U.shape)
            self.assertEqual(v.shape, V.shape)
            self.assertEqual(s, S)

            if density == 1:
                # actual_rank is known only for dense inputs
                #
                # check if pairs (u, U) and (v, V) span the same
                # subspaces, respectively
                u, s, v = u[..., :actual_rank], s[..., :actual_rank], v[..., :actual_rank]
                U, S, V = U[..., :actual_rank], S[..., :actual_rank], V[..., :actual_rank]
                self.assertEqual(u.mT.matmul(U).det().abs(), torch.ones(batches, device=device, dtype=dtype))
                self.assertEqual(v.mT.matmul(V).det().abs(), torch.ones(batches, device=device, dtype=dtype))

        all_batches = [(), (1,), (3,), (2, 3)]
        for actual_rank, size, all_batches in [
                (2, (17, 4), all_batches),
                (4, (17, 4), all_batches),
                (4, (17, 17), all_batches),
                (10, (100, 40), all_batches),
                (7, (1000, 1000), [()]),
        ]:
            # dense input
            for batches in all_batches:
                run_subtest(actual_rank, size, batches, device, torch.svd_lowrank)
                if size != size[::-1]:
                    run_subtest(actual_rank, size[::-1], batches, device, torch.svd_lowrank)

        # sparse input
        for size in [(17, 4), (4, 17), (17, 17), (100, 40), (40, 100), (1000, 1000)]:
            for density in [0.005, 0.1]:
                run_subtest(None, size, (), device, torch.svd_lowrank, density=density)

        # jitting support
        jitted = torch.jit.script(torch.svd_lowrank)
        actual_rank, size, batches = 2, (17, 4), ()
        run_subtest(actual_rank, size, batches, device, jitted)
//  [python]
def _r4(X, Y, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    return _r1(torch.svd(X).U, torch.svd(Y).U)
//  [python]
def _r4(X, Y, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    return _r1(torch.linalg.svd(X).U, torch.linalg.svd(Y).U)
//  [python]
def init_projection_matrix(self, x):
        # Set if using projection matrix
        cfg = self.cfg
        self.use_projection_matrix = cfg["use_projection_matrix"]

        if self.use_projection_matrix:
            self.compressed_dim = TensorList([cfg["compressed_dim"]])
            proj_init_method = cfg["proj_init_method"]
            if proj_init_method == 'pca':
                x_mat = TensorList([
                    e.permute(1, 0, 2, 3).reshape(e.shape[1], -1).clone()
                    for e in x
                ])
                x_mat -= x_mat.mean(dim=1, keepdim=True)
                cov_x = x_mat @ x_mat.t()
                self.projection_matrix = TensorList([
                    None if cdim is None else torch.svd(C)[0]
                    [:, :cdim].t().unsqueeze(-1).unsqueeze(-1).clone()
                    for C, cdim in zip(cov_x, self.compressed_dim)
                ])
            elif proj_init_method == 'randn':
                self.projection_matrix = TensorList([
                    None if cdim is None else ex.new_zeros(
                        cdim, ex.shape[1], 1, 1).normal_(
                            0, 1 / math.sqrt(ex.shape[1]))
                    for ex, cdim in zip(x, self.compressed_dim)
                ])
        else:
            self.compressed_dim = x.size(1)
            self.projection_matrix = TensorList([None] * len(x))
//  [python]
def init_projection_matrix(self, x):
        # Set if using projection matrix
        cfg = self.cfg
        self.use_projection_matrix = cfg["use_projection_matrix"]

        if self.use_projection_matrix:
            self.compressed_dim = TensorList([cfg["compressed_dim"]])
            proj_init_method = cfg["proj_init_method"]
            if proj_init_method == 'pca':
                x_mat = TensorList([
                    e.permute(1, 0, 2, 3).reshape(e.shape[1], -1).clone()
                    for e in x
                ])
                x_mat -= x_mat.mean(dim=1, keepdim=True)
                cov_x = x_mat @ x_mat.t()
                self.projection_matrix = TensorList([
                    None if cdim is None else torch.linalg.svd(C)[0]
                    [:, :cdim].t().unsqueeze(-1).unsqueeze(-1).clone()
                    for C, cdim in zip(cov_x, self.compressed_dim)
                ])
            elif proj_init_method == 'randn':
                self.projection_matrix = TensorList([
                    None if cdim is None else ex.new_zeros(
                        cdim, ex.shape[1], 1, 1).normal_(
                            0, 1 / math.sqrt(ex.shape[1]))
                    for ex, cdim in zip(x, self.compressed_dim)
                ])
        else:
            self.compressed_dim = x.size(1)
            self.projection_matrix = TensorList([None] * len(x))
//  [python]
def _r3(X, Y, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    UX, _, VX = torch.svd(X)
    UY, _, VY = torch.svd(Y)
    return _r1(UX @ VX.T, UY @ VY.T)
//  [python]
def _r3(X, Y, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    UX, S, VX = torch.linalg.svd(X)
    UY, S, VY = torch.linalg.svd(Y)
    return _r1(UX @ VX.T, UY @ VY.T)
//  [python]
def _rG(X, Y, n_components: Optional[int] = None, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    ncomp1 = _matrix_rank(X) if n_components is None else n_components
    ncomp2 = _matrix_rank(Y) if n_components is None else n_components
    UX = torch.svd(X).U[:, :ncomp1]
    UY = torch.svd(Y).U[:, :ncomp2]
    return _r1(UX.T @ UX, UY.T @ UY)
//  [python]
def _rG(X, Y, n_components: Optional[int] = None, center: bool = True):
    if center:
        X = X - X.mean(0)
        Y = Y - Y.mean(0)
    ncomp1 = _matrix_rank(X) if n_components is None else n_components
    ncomp2 = _matrix_rank(Y) if n_components is None else n_components
    UX = torch.linalg.svd(X).U[:, :ncomp1]
    UY = torch.linalg.svd(Y).U[:, :ncomp2]
    return _r1(UX.T @ UX, UY.T @ UY)
//  [python]
def forward(self, feature):
        result = 0
        u, s, v = torch.svd(feature.t())
        num = s.size(0)
        for i in range(self.k):
            result += torch.pow(s[num-1-i], 2)
        return result
//  [python]
def forward(self, feature):
        result = 0
        u, s, v = torch.linalg.svd(feature.t())
        num = s.size(0)
        for i in range(self.k):
            result += torch.pow(s[num-1-i], 2)
        return result
//  [python]
def _svd(self,
        points1: torch.Tensor,
        points2: torch.Tensor
    ) -> torch.Tensor:
        '''
            Computes a similarity transform (sR, t) that takes
            a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
            where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
            i.e. solves the orthogonal Procrutes problem.
        '''
        #NOTE: adapted from https://gist.github.com/mkocabas/54ea2ff3b03260e3fedf8ad22536f427
        b, _, c = points1.shape
        # 1. Remove mean.
        points1 = torch.transpose(points1, -2, -1)
        points2 = torch.transpose(points2, -2, -1)
        centroid1 = points1.mean(dim=-1, keepdims=True)
        centroid2 = points1.mean(dim=-1, keepdims=True)
        centered1 = points1 - centroid1
        centered2 = points2 - centroid2
        # 2. Compute variance of X1 used for scale.
        variance = torch.sum(centered1 ** 2, dim=[1, 2])
        # 3. The outer product of X1 and X2.    
        K = centered1 @ torch.transpose(centered2, -2, -1)
        # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are singular vectors of K.
        U, s, V = torch.svd(K)
        # Construct Z that fixes the orientation of R to get det(R)=1.
        Z = torch.eye(c).to(U).unsqueeze(0).repeat(b, 1, 1)
        Z[:,-1, -1] *= torch.sign(torch.det(U @ torch.transpose(V, -2, -1)))
        # Construct R.
        rotation = V @ (Z @ torch.transpose(U, -2, -1))
        # 5. Recover scale.
        scale = torch.cat([torch.trace(x).unsqueeze(0) for x in (rotation @ K)]) / variance
        # 6. Recover translation.
        scale = scale.unsqueeze(-1).unsqueeze(-1)
        translation = centroid2 - (scale * (rotation @ centroid1))
        return rotation, translation, scale
//  [python]
def _svd(self,
        points1: torch.Tensor,
        points2: torch.Tensor
    ) -> torch.Tensor:
        '''
            Computes a similarity transform (sR, t) that takes
            a set of 3D points S1 (3 x N) closest to a set of 3D points S2,
            where R is an 3x3 rotation matrix, t 3x1 translation, s scale.
            i.e. solves the orthogonal Procrutes problem.
        '''
        #NOTE: adapted from https://gist.github.com/mkocabas/54ea2ff3b03260e3fedf8ad22536f427
        b, _, c = points1.shape
        # 1. Remove mean.
        points1 = torch.transpose(points1, -2, -1)
        points2 = torch.transpose(points2, -2, -1)
        centroid1 = points1.mean(dim=-1, keepdims=True)
        centroid2 = points1.mean(dim=-1, keepdims=True)
        centered1 = points1 - centroid1
        centered2 = points2 - centroid2
        # 2. Compute variance of X1 used for scale.
        variance = torch.sum(centered1 ** 2, dim=[1, 2])
        # 3. The outer product of X1 and X2.    
        K = centered1 @ torch.transpose(centered2, -2, -1)
        # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are singular vectors of K.
        U, s, V = torch.linalg.svd(K)
        # Construct Z that fixes the orientation of R to get det(R)=1.
        Z = torch.eye(c).to(U).unsqueeze(0).repeat(b, 1, 1)
        Z[:,-1, -1] *= torch.sign(torch.det(U @ torch.transpose(V, -2, -1)))
        # Construct R.
        rotation = V @ (Z @ torch.transpose(U, -2, -1))
        # 5. Recover scale.
        scale = torch.cat([torch.trace(x).unsqueeze(0) for x in (rotation @ K)]) / variance
        # 6. Recover translation.
        scale = scale.unsqueeze(-1).unsqueeze(-1)
        translation = centroid2 - (scale * (rotation @ centroid1))
        return rotation, translation, scale
//  [python]
def _matrix_power(matrix, power):
    # use CPU for svd for speed up
    device = matrix.device
    matrix = matrix.cpu()
    u, s, v = torch.svd(matrix)
    return (u @ s.pow_(power).diag() @ v.t()).to(device)
//  [python]
def _matrix_power(matrix, power):
    # use CPU for svd for speed up
    device = matrix.device
    matrix = matrix.cpu()
    u, s, v = torch.linalg.svd(matrix)
    return (u @ s.pow_(power).diag() @ v.t()).to(device)
//  [python]
def gmm_register(pi_s, mu_s, mu_t, sigma_t):
	'''
	Inputs:
		pi: B x J
		mu: B x J x 3
		sigma: B x J x 3 x 3
	'''
	c_s = pi_s.unsqueeze(1) @ mu_s
	c_t = pi_s.unsqueeze(1) @ mu_t
	Ms = torch.sum((pi_s.unsqueeze(2) * (mu_s - c_s)).unsqueeze(3) @
				   (mu_t - c_t).unsqueeze(2) @ sigma_t.inverse(), dim=1)
	U, _, V = torch.svd(Ms.cpu())
	U = U.cuda() if torch.cuda.is_available() else U
	V = V.cuda() if torch.cuda.is_available() else V
	S = torch.eye(3).unsqueeze(0).repeat(U.shape[0], 1, 1).to(U.device)
	S[:, 2, 2] = torch.det(V @ U.transpose(1, 2))
	R = V @ S @ U.transpose(1, 2)
	t = c_t.transpose(1, 2) - R @ c_s.transpose(1, 2)
	bot_row = torch.Tensor([[[0, 0, 0, 1]]]).repeat(R.shape[0], 1, 1).to(R.device)
	T = torch.cat([torch.cat([R, t], dim=2), bot_row], dim=1)
	return T
//  [python]
def gmm_register(pi_s, mu_s, mu_t, sigma_t):
	'''
	Inputs:
		pi: B x J
		mu: B x J x 3
		sigma: B x J x 3 x 3
	'''
	c_s = pi_s.unsqueeze(1) @ mu_s
	c_t = pi_s.unsqueeze(1) @ mu_t
	Ms = torch.sum((pi_s.unsqueeze(2) * (mu_s - c_s)).unsqueeze(3) @
				   (mu_t - c_t).unsqueeze(2) @ sigma_t.inverse(), dim=1)
	U, S, V = torch.linalg.svd(Ms.cpu())
	U = U.cuda() if torch.cuda.is_available() else U
	V = V.cuda() if torch.cuda.is_available() else V
	S = torch.eye(3).unsqueeze(0).repeat(U.shape[0], 1, 1).to(U.device)
	S[:, 2, 2] = torch.det(V @ U.transpose(1, 2))
	R = V @ S @ U.transpose(1, 2)
	t = c_t.transpose(1, 2) - R @ c_s.transpose(1, 2)
	bot_row = torch.Tensor([[[0, 0, 0, 1]]]).repeat(R.shape[0], 1, 1).to(R.device)
	T = torch.cat([torch.cat([R, t], dim=2), bot_row], dim=1)
	return T
//  [python]
def forward(self, module_input):
        # Before converting the output rotation matrices of the VAE to
        # axis-angle representation, we first need to make them in to valid
        # rotation matrices
        with torch.no_grad():
            # TODO: Replace with Batch SVD once merged
            # Iterate over the batch dimension and compute the SVD
            svd_input = module_input.detach().cpu()
            #  svd_input = output
            norm_rotation = torch.zeros_like(svd_input)
            for bidx in range(module_input.shape[0]):
                U, _, V = torch.svd(svd_input[bidx])

                # Multiply the U, V matrices to get the closest orthonormal
                # matrix
                norm_rotation[bidx] = torch.matmul(U, V.t())
            norm_rotation = norm_rotation.to(module_input.device)

        # torch.svd supports backprop only for full-rank matrices.
        # The output is calculated as the valid rotation matrix plus the
        # output minus the detached output. If one writes down the
        # computational graph for this operation, it will become clear the
        # output is the desired valid rotation matrix, while for the
        # backward pass gradients are propagated only to the original
        # matrix
        # Source: PyTorch Gumbel-Softmax hard sampling
        # https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#gumbel_softmax
        correct_rot = norm_rotation - module_input.detach() + module_input
        return correct_rot
//  [python]
def forward(self, module_input):
        # Before converting the output rotation matrices of the VAE to
        # axis-angle representation, we first need to make them in to valid
        # rotation matrices
        with torch.no_grad():
            # TODO: Replace with Batch SVD once merged
            # Iterate over the batch dimension and compute the SVD
            svd_input = module_input.detach().cpu()
            #  svd_input = output
            norm_rotation = torch.zeros_like(svd_input)
            for bidx in range(module_input.shape[0]):
                U, S, V = torch.linalg.svd(svd_input[bidx])

                # Multiply the U, V matrices to get the closest orthonormal
                # matrix
                norm_rotation[bidx] = torch.matmul(U, V.t())
            norm_rotation = norm_rotation.to(module_input.device)

        # torch.svd supports backprop only for full-rank matrices.
        # The output is calculated as the valid rotation matrix plus the
        # output minus the detached output. If one writes down the
        # computational graph for this operation, it will become clear the
        # output is the desired valid rotation matrix, while for the
        # backward pass gradients are propagated only to the original
        # matrix
        # Source: PyTorch Gumbel-Softmax hard sampling
        # https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#gumbel_softmax
        correct_rot = norm_rotation - module_input.detach() + module_input
        return correct_rot
//  [python]
def svd(x):  # pragma: no cover
    # inspired by
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    # prolonged here:
    if x.dim() == 2:
        # 17 milliseconds on my mac to check that condition, that is low overhead
        result = torch.svd(x)
    else:
        batches = x.shape[:-2]
        other = x.shape[-2:]
        flat = x.view((-1,) + other)
        slices = flat.unbind(0)
        U, D, V = [], [], []
        # I wish I had a parallel_for
        for i in range(flat.shape[0]):
            u, d, v = torch.svd(slices[i])
            U += [u]
            D += [d]
            V += [v]
        U = torch.stack(U).view(batches + U[0].shape)
        D = torch.stack(D).view(batches + D[0].shape)
        V = torch.stack(V).view(batches + V[0].shape)
        result = U, D, V
    return result
//  [python]
def svd(x):  # pragma: no cover
    # inspired by
    # https://discuss.pytorch.org/t/multidimensional-svd/4366/2
    # prolonged here:
    if x.dim() == 2:
        # 17 milliseconds on my mac to check that condition, that is low overhead
        result = torch.linalg.svd(x)
    else:
        batches = x.shape[:-2]
        other = x.shape[-2:]
        flat = x.view((-1,) + other)
        slices = flat.unbind(0)
        U, D, V = [], [], []
        # I wish I had a parallel_for
        for i in range(flat.shape[0]):
            u, d, v = torch.linalg.svd(slices[i])
            U += [u]
            D += [d]
            V += [v]
        U = torch.stack(U).view(batches + U[0].shape)
        D = torch.stack(D).view(batches + D[0].shape)
        V = torch.stack(V).view(batches + V[0].shape)
        result = U, D, V
    return result