jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_perplexity(self) -> None:

        input = torch.tensor([[[0.3659, 0.7025, 0.3104]], [[0.0097, 0.6577, 0.1947]]])
        target = torch.tensor([[2], [1]])
        t.assert_close(
            perplexity(input, target),
            torch.tensor(2.759287357330, dtype=torch.float64),
        )

        input = torch.tensor(
            [
                [
                    [0.3633, 0.7777, 0.3111, 0.1134],
                    [0.5543, 0.4387, 0.8176, 0.4565],
                    [0.1324, 0.1465, 0.2822, 0.6654],
                ],
                [
                    [0.0956, 0.6523, 0.8914, 0.3653],
                    [0.3289, 0.7254, 0.8121, 0.4196],
                    [0.3233, 0.7775, 0.8731, 0.4001],
                ],
            ]
        )
        target = torch.tensor([[2, 1, 3], [1, 0, 1]])
        t.assert_close(
            perplexity(input, target),
            torch.tensor(3.944433212280, dtype=torch.float64),
        )
//  [python]
def test_perplexity(self) -> None:

        input = torch.tensor([[[0.3659, 0.7025, 0.3104]], [[0.0097, 0.6577, 0.1947]]])
        target = torch.tensor([[2], [1]])
        torch.testing.assert_close(
            perplexity(input, target),
            torch.tensor(2.759287357330, dtype=torch.float64),
        )

        input = torch.tensor(
            [
                [
                    [0.3633, 0.7777, 0.3111, 0.1134],
                    [0.5543, 0.4387, 0.8176, 0.4565],
                    [0.1324, 0.1465, 0.2822, 0.6654],
                ],
                [
                    [0.0956, 0.6523, 0.8914, 0.3653],
                    [0.3289, 0.7254, 0.8121, 0.4196],
                    [0.3233, 0.7775, 0.8731, 0.4001],
                ],
            ]
        )
        target = torch.tensor([[2, 1, 3], [1, 0, 1]])
        torch.testing.assert_close(
            perplexity(input, target),
            torch.tensor(3.944433212280, dtype=torch.float64),
        )
//  [python]
def test_perplexity_with_ignore_index(self) -> None:

        input = torch.tensor([[[0.3659, 0.7025, 0.3104]], [[0.0097, 0.6577, 0.1947]]])
        target = torch.tensor([[2], [100]])
        torch.testing.assert_allclose(
            perplexity(input, target, ignore_index=100),
            torch.tensor(3.537154912949, dtype=torch.float64),
        )

        input = torch.tensor(
            [
                [
                    [0.3633, 0.7777, 0.3111, 0.1134],
                    [0.5543, 0.4387, 0.8176, 0.4565],
                    [0.1324, 0.1465, 0.2822, 0.6654],
                ],
                [
                    [0.0956, 0.6523, 0.8914, 0.3653],
                    [0.3289, 0.7254, 0.8121, 0.4196],
                    [0.3233, 0.7775, 0.8731, 0.4001],
                ],
            ]
        )
        target = torch.tensor([[2, 1, 3], [1, 0, 1]])
        torch.testing.assert_allclose(
            perplexity(input, target, ignore_index=1),
            torch.tensor(4.053068637848, dtype=torch.float64),
        )
//  [python]
def test_perplexity_with_ignore_index(self) -> None:

        input = torch.tensor([[[0.3659, 0.7025, 0.3104]], [[0.0097, 0.6577, 0.1947]]])
        target = torch.tensor([[2], [100]])
        torch.testing.assert_close(
            perplexity(input, target, ignore_index=100),
            torch.tensor(3.537154912949, dtype=torch.float64),
        )

        input = torch.tensor(
            [
                [
                    [0.3633, 0.7777, 0.3111, 0.1134],
                    [0.5543, 0.4387, 0.8176, 0.4565],
                    [0.1324, 0.1465, 0.2822, 0.6654],
                ],
                [
                    [0.0956, 0.6523, 0.8914, 0.3653],
                    [0.3289, 0.7254, 0.8121, 0.4196],
                    [0.3233, 0.7775, 0.8731, 0.4001],
                ],
            ]
        )
        target = torch.tensor([[2, 1, 3], [1, 0, 1]])
        torch.testing.assert_close(
            perplexity(input, target, ignore_index=1),
            torch.tensor(4.053068637848, dtype=torch.float64),
        )
//  [python]
def test_model_inference_batched(self):
        # Load sample checkpoint from Tortoise TTS
        model = UnivNetModel.from_pretrained("dg845/univnet-dev")
        model.eval().to(torch_device)

        # Get batched noise and spectrogram inputs.
        input_speech = self.get_inputs(torch_device, num_samples=3)

        with torch.no_grad():
            waveform = model(**input_speech)[0]
        waveform = waveform.cpu()

        waveform_mean = torch.mean(waveform)
        waveform_stddev = torch.std(waveform)
        waveform_slice = waveform[-1, -9:].flatten()

        EXPECTED_MEAN = torch.tensor(-0.19989729)
        EXPECTED_STDDEV = torch.tensor(0.35230172)
        EXPECTED_SLICE = torch.tensor([-0.3408, -0.6045, -0.5052, 0.1160, -0.1556, -0.0405, -0.3024, -0.5290, -0.5019])

        torch.testing.assert_allclose(waveform_mean, EXPECTED_MEAN, atol=1e-4, rtol=1e-5)
        torch.testing.assert_allclose(waveform_stddev, EXPECTED_STDDEV, atol=1e-4, rtol=1e-5)
        torch.testing.assert_allclose(waveform_slice, EXPECTED_SLICE, atol=5e-4, rtol=1e-5)
//  [python]
def test_model_inference_batched(self):
        # Load sample checkpoint from Tortoise TTS
        model = UnivNetModel.from_pretrained("dg845/univnet-dev")
        model.eval().to(torch_device)

        # Get batched noise and spectrogram inputs.
        input_speech = self.get_inputs(torch_device, num_samples=3)

        with torch.no_grad():
            waveform = model(**input_speech)[0]
        waveform = waveform.cpu()

        waveform_mean = torch.mean(waveform)
        waveform_stddev = torch.std(waveform)
        waveform_slice = waveform[-1, -9:].flatten()

        EXPECTED_MEAN = torch.tensor(-0.19989729)
        EXPECTED_STDDEV = torch.tensor(0.35230172)
        EXPECTED_SLICE = torch.tensor([-0.3408, -0.6045, -0.5052, 0.1160, -0.1556, -0.0405, -0.3024, -0.5290, -0.5019])

        torch.testing.assert_close(waveform_mean, EXPECTED_MEAN, atol=1e-4, rtol=1e-5)
        torch.testing.assert_close(waveform_stddev, EXPECTED_STDDEV, atol=1e-4, rtol=1e-5)
        torch.testing.assert_close(waveform_slice, EXPECTED_SLICE, atol=5e-4, rtol=1e-5)
//  [python]
def test_model_inference_unbatched(self):
        # Load sample checkpoint from Tortoise TTS
        model = UnivNetModel.from_pretrained("dg845/univnet-dev")
        model.eval().to(torch_device)

        # Get unbatched noise and spectrogram inputs.
        input_speech = self.get_inputs(torch_device, num_samples=1)

        with torch.no_grad():
            waveform = model(**input_speech)[0]
        waveform = waveform.cpu()

        waveform_mean = torch.mean(waveform)
        waveform_stddev = torch.std(waveform)
        waveform_slice = waveform[-1, -9:].flatten()

        EXPECTED_MEAN = torch.tensor(-0.22895093)
        EXPECTED_STDDEV = torch.tensor(0.33986747)
        EXPECTED_SLICE = torch.tensor([-0.3276, -0.5504, -0.3484, 0.3574, -0.0373, -0.1826, -0.4880, -0.6431, -0.5162])

        torch.testing.assert_allclose(waveform_mean, EXPECTED_MEAN, atol=1e-4, rtol=1e-5)
        torch.testing.assert_allclose(waveform_stddev, EXPECTED_STDDEV, atol=1e-4, rtol=1e-5)
        torch.testing.assert_allclose(waveform_slice, EXPECTED_SLICE, atol=1e-3, rtol=1e-5)
//  [python]
def test_model_inference_unbatched(self):
        # Load sample checkpoint from Tortoise TTS
        model = UnivNetModel.from_pretrained("dg845/univnet-dev")
        model.eval().to(torch_device)

        # Get unbatched noise and spectrogram inputs.
        input_speech = self.get_inputs(torch_device, num_samples=1)

        with torch.no_grad():
            waveform = model(**input_speech)[0]
        waveform = waveform.cpu()

        waveform_mean = torch.mean(waveform)
        waveform_stddev = torch.std(waveform)
        waveform_slice = waveform[-1, -9:].flatten()

        EXPECTED_MEAN = torch.tensor(-0.22895093)
        EXPECTED_STDDEV = torch.tensor(0.33986747)
        EXPECTED_SLICE = torch.tensor([-0.3276, -0.5504, -0.3484, 0.3574, -0.0373, -0.1826, -0.4880, -0.6431, -0.5162])

        torch.testing.assert_close(waveform_mean, EXPECTED_MEAN, atol=1e-4, rtol=1e-5)
        torch.testing.assert_close(waveform_stddev, EXPECTED_STDDEV, atol=1e-4, rtol=1e-5)
        torch.testing.assert_close(waveform_slice, EXPECTED_SLICE, atol=1e-3, rtol=1e-5)
//  [python]
def test_lazy_load_subclass(lit_llama):
    import lit_llama.utils

    with tempfile.TemporaryDirectory() as tmpdirname:
        path = pathlib.Path(tmpdirname)
        fn = str(path / "test.pt")
        t = torch.randn(2, 3)[:, 1:]
        sd = {
            1: t,
            2: torch.nn.Parameter(t),
            3: torch.Tensor._make_subclass(ATensor, t),
        }
        torch.save(sd, fn)
        with lit_llama.utils.lazy_load(fn) as sd_lazy:
            for k in sd.keys():
                actual = sd_lazy[k]
                expected = sd[k]
                torch.testing.assert_allclose(actual._load_tensor(), expected)
//  [python]
def test_lazy_load_subclass(lit_llama):
    import lit_llama.utils

    with tempfile.TemporaryDirectory() as tmpdirname:
        path = pathlib.Path(tmpdirname)
        fn = str(path / "test.pt")
        t = torch.randn(2, 3)[:, 1:]
        sd = {
            1: t,
            2: torch.nn.Parameter(t),
            3: torch.Tensor._make_subclass(ATensor, t),
        }
        torch.save(sd, fn)
        with lit_llama.utils.lazy_load(fn) as sd_lazy:
            for k in sd.keys():
                actual = sd_lazy[k]
                expected = sd[k]
                torch.testing.assert_close(actual._load_tensor(), expected)
//  [python]
def test_integration(self):
        feature_extractor = UnivNetFeatureExtractor.from_pretrained("dg845/univnet-dev")
        model = UnivNetModel.from_pretrained("dg845/univnet-dev")
        model.eval().to(torch_device)

        audio, sr = self._load_datasamples(1, sampling_rate=feature_extractor.sampling_rate)

        input_features = feature_extractor(audio, sampling_rate=sr[0], return_tensors="pt").input_features
        input_features = input_features.to(device=torch_device)

        input_speech = self.get_inputs(torch_device, num_samples=1, noise_length=input_features.shape[1])
        input_speech["input_features"] = input_features

        with torch.no_grad():
            waveform = model(**input_speech)[0]
        waveform = waveform.cpu()

        waveform_mean = torch.mean(waveform)
        waveform_stddev = torch.std(waveform)
        waveform_slice = waveform[-1, -9:].flatten()

        EXPECTED_MEAN = torch.tensor(0.00051374)
        EXPECTED_STDDEV = torch.tensor(0.058105603)
        # fmt: off
        EXPECTED_SLICE = torch.tensor([-4.3934e-04, -1.8203e-04, -3.3033e-04, -3.8716e-04, -1.6125e-04, 3.5389e-06, -3.3149e-04, -3.7613e-04, -2.3331e-04])
        # fmt: on

        torch.testing.assert_allclose(waveform_mean, EXPECTED_MEAN, atol=5e-6, rtol=1e-5)
        torch.testing.assert_allclose(waveform_stddev, EXPECTED_STDDEV, atol=1e-4, rtol=1e-5)
        torch.testing.assert_allclose(waveform_slice, EXPECTED_SLICE, atol=5e-6, rtol=1e-5)
//  [python]
def test_integration(self):
        feature_extractor = UnivNetFeatureExtractor.from_pretrained("dg845/univnet-dev")
        model = UnivNetModel.from_pretrained("dg845/univnet-dev")
        model.eval().to(torch_device)

        audio, sr = self._load_datasamples(1, sampling_rate=feature_extractor.sampling_rate)

        input_features = feature_extractor(audio, sampling_rate=sr[0], return_tensors="pt").input_features
        input_features = input_features.to(device=torch_device)

        input_speech = self.get_inputs(torch_device, num_samples=1, noise_length=input_features.shape[1])
        input_speech["input_features"] = input_features

        with torch.no_grad():
            waveform = model(**input_speech)[0]
        waveform = waveform.cpu()

        waveform_mean = torch.mean(waveform)
        waveform_stddev = torch.std(waveform)
        waveform_slice = waveform[-1, -9:].flatten()

        EXPECTED_MEAN = torch.tensor(0.00051374)
        EXPECTED_STDDEV = torch.tensor(0.058105603)
        # fmt: off
        EXPECTED_SLICE = torch.tensor([-4.3934e-04, -1.8203e-04, -3.3033e-04, -3.8716e-04, -1.6125e-04, 3.5389e-06, -3.3149e-04, -3.7613e-04, -2.3331e-04])
        # fmt: on

        torch.testing.assert_close(waveform_mean, EXPECTED_MEAN, atol=5e-6, rtol=1e-5)
        torch.testing.assert_close(waveform_stddev, EXPECTED_STDDEV, atol=1e-4, rtol=1e-5)
        torch.testing.assert_close(waveform_slice, EXPECTED_SLICE, atol=5e-6, rtol=1e-5)
//  [python]
def test_incremental_write(tmp_path, lit_llama):
    import lit_llama.utils

    sd = {str(k): torch.randn(5, 10) for k in range(3)}
    sd_expected = {k: v.clone() for k, v in sd.items()}
    fn = str(tmp_path / "test.pt")
    with lit_llama.utils.incremental_save(fn) as f:
        sd["0"] = f.store_early(sd["0"])
        sd["2"] = f.store_early(sd["2"])
        f.save(sd)
    sd_actual = torch.load(fn)
    assert sd_actual.keys() == sd_expected.keys()
    for k, v_expected in sd_expected.items():
        v_actual = sd_actual[k]
        torch.testing.assert_allclose(v_expected, v_actual)
//  [python]
def test_incremental_write(tmp_path, lit_llama):
    import lit_llama.utils

    sd = {str(k): torch.randn(5, 10) for k in range(3)}
    sd_expected = {k: v.clone() for k, v in sd.items()}
    fn = str(tmp_path / "test.pt")
    with lit_llama.utils.incremental_save(fn) as f:
        sd["0"] = f.store_early(sd["0"])
        sd["2"] = f.store_early(sd["2"])
        f.save(sd)
    sd_actual = torch.load(fn)
    assert sd_actual.keys() == sd_expected.keys()
    for k, v_expected in sd_expected.items():
        v_actual = sd_actual[k]
        torch.testing.assert_close(v_expected, v_actual)
//  [python]
def svd(feat, iden=False, device='cpu'):
    size = feat.size()
    mean = torch.mean(feat, 1)
    mean = mean.unsqueeze(1).expand_as(feat)
    _feat = feat.clone()
    _feat -= mean
    if size[1] > 1:
        conv = torch.mm(_feat, _feat.t()).div(size[1] - 1)
    else:
        conv = torch.mm(_feat, _feat.t())
    if iden:
        conv += torch.eye(size[0]).to(device)
    u, e, v = torch.svd(conv, some=False)
    return u, e, v
//  [python]
def svd(feat, iden=False, device='cpu'):
    size = feat.size()
    mean = torch.mean(feat, 1)
    mean = mean.unsqueeze(1).expand_as(feat)
    _feat = feat.clone()
    _feat -= mean
    if size[1] > 1:
        conv = torch.mm(_feat, _feat.t()).div(size[1] - 1)
    else:
        conv = torch.mm(_feat, _feat.t())
    if iden:
        conv += torch.eye(size[0]).to(device)
    u, e, v = torch.linalg.svd(conv)
    return u, e, v
//  [python]
def test_block_ptr_matmul_no_scf(shape, num_warps):
    capability = torch.cuda.get_device_capability()
    if torch.version.hip is None and capability[0] >= 9:
        pytest.skip("Hopper support is working in progress")

    m, n, k = shape
    a = torch.randn((m, k), device="cuda", dtype=torch.float16)
    b = torch.randn((k, n), device="cuda", dtype=torch.float16)
    c = torch.empty((m, n), device="cuda", dtype=torch.float32)

    grid = lambda META: (1, )
    matmul_no_scf_with_advance_kernel[grid](
        a_ptr=a, b_ptr=b, c_ptr=c,  #
        M=m, N=n, K=k,  #
        stride_am=a.stride(0), stride_ak=a.stride(1),  #
        stride_bk=b.stride(0), stride_bn=b.stride(1),  #
        stride_cm=c.stride(0), stride_cn=c.stride(1),  #
        BLOCK_M=m, BLOCK_N=n, BLOCK_K=k,  #
        num_warps=num_warps)
    golden = torch.matmul(a, b)
    torch.testing.assert_allclose(c, golden, check_dtype=False)
//  [python]
def test_block_ptr_matmul_no_scf(shape, num_warps):
    capability = torch.cuda.get_device_capability()
    if torch.version.hip is None and capability[0] >= 9:
        pytest.skip("Hopper support is working in progress")

    m, n, k = shape
    a = torch.randn((m, k), device="cuda", dtype=torch.float16)
    b = torch.randn((k, n), device="cuda", dtype=torch.float16)
    c = torch.empty((m, n), device="cuda", dtype=torch.float32)

    grid = lambda META: (1, )
    matmul_no_scf_with_advance_kernel[grid](
        a_ptr=a, b_ptr=b, c_ptr=c,  #
        M=m, N=n, K=k,  #
        stride_am=a.stride(0), stride_ak=a.stride(1),  #
        stride_bk=b.stride(0), stride_bn=b.stride(1),  #
        stride_cm=c.stride(0), stride_cn=c.stride(1),  #
        BLOCK_M=m, BLOCK_N=n, BLOCK_K=k,  #
        num_warps=num_warps)
    golden = torch.matmul(a, b)
    torch.testing.assert_close(c, golden, check_dtype=False)
//  [python]
def kabsch_torch(A, B, C):
    A=A.double()
    B=B.double()
    C=C.double()
    a_mean = A.mean(dim=0, keepdims=True)
    b_mean = B.mean(dim=0, keepdims=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = torch.matmul(A_c.transpose(0,1), B_c)  # [B, 3, 3]
    U, S, V = torch.svd(H)
    # Rotation matrix
    R = torch.matmul(V, U.transpose(0,1))  # [B, 3, 3]
    # Translation vector
    t = b_mean - torch.matmul(R, a_mean.transpose(0,1)).transpose(0,1)
    C_aligned = torch.matmul(R, C.transpose(0,1)).transpose(0,1) + t
    return C_aligned, R, t
//  [python]
def kabsch_torch(A, B, C):
    A=A.double()
    B=B.double()
    C=C.double()
    a_mean = A.mean(dim=0, keepdims=True)
    b_mean = B.mean(dim=0, keepdims=True)
    A_c = A - a_mean
    B_c = B - b_mean
    # Covariance matrix
    H = torch.matmul(A_c.transpose(0,1), B_c)  # [B, 3, 3]
    U, S, V = torch.linalg.svd(H)
    # Rotation matrix
    R = torch.matmul(V, U.transpose(0,1))  # [B, 3, 3]
    # Translation vector
    t = b_mean - torch.matmul(R, a_mean.transpose(0,1)).transpose(0,1)
    C_aligned = torch.matmul(R, C.transpose(0,1)).transpose(0,1) + t
    return C_aligned, R, t
//  [python]
def whiten_tensor_svd(X):

    X_c = X - X.mean()
    Sigma = X_c.transpose(0,1).matmul(X_c) / X_c.shape[0]
    Sigma = Sigma.cpu()
    # try:
    U, Lambda, _ = torch.svd(Sigma)
    # except:                     # torch.svd may have convergence issues for GPU and CPU.
    # U, Lambda, _ = torch.svd(Sigma + 1e-4*Sigma.mean()*torch.randn_like(Sigma))
    # U, Lambda, _ = torch.svd()
    U = U.to(X.device)
    Lambda = Lambda.to(X.device)
    W = U.matmul(torch.diag(1.0/torch.sqrt(Lambda + 1e-5)).matmul(U.transpose(0,1)))
    return X_c.matmul(W.transpose(0,1))
//  [python]
def whiten_tensor_svd(X):

    X_c = X - X.mean()
    Sigma = X_c.transpose(0,1).matmul(X_c) / X_c.shape[0]
    Sigma = Sigma.cpu()
    U, S, V = torch.linalg.svd(Sigma)
    U = U.to(X.device)
    S = S.to(X.device)
    W = U.matmul(torch.diag(1.0/torch.sqrt(S + 1e-5)).matmul(U.transpose(0,1)))
    return X_c.matmul(W.transpose(0,1))
//  [python]
def _mat_sqrt(x):
    U, D, V = torch.svd(x)
    return torch.mm(torch.mm(U, D.pow(0.5).diag()), V.t())
//  [python]
def _mat_sqrt(x):
    U, S, V = torch.linalg.svd(x)
    return torch.mm(torch.mm(U, S.pow(0.5).diag()), V.t())
//  [python]
def __init__(self, dims_in, correction_interval=256, clamp=5.):
        super().__init__()
        self.width = dims_in[0][0]
        self.clamp = clamp

        self.correction_interval = correction_interval
        self.back_counter = np.random.randint(0, correction_interval) // 2

        self.weights = torch.randn(self.width, self.width)
        self.weights = self.weights + self.weights.t()
        self.weights, S, V = torch.svd(self.weights)

        self.weights = nn.Parameter(self.weights)

        self.bias = nn.Parameter(0.05 * torch.randn(self.width))
        self.scaling = nn.Parameter(0.02 * torch.randn(self.width))

        self.register_backward_hook(correct_weights)
//  [python]
def __init__(self, dims_in, correction_interval=256, clamp=5.):
        super().__init__()
        self.width = dims_in[0][0]
        self.clamp = clamp

        self.correction_interval = correction_interval
        self.back_counter = np.random.randint(0, correction_interval) // 2

        self.weights = torch.randn(self.width, self.width)
        self.weights = self.weights + self.weights.t()
        self.weights, S, V = torch.linalg.svd(self.weights)

        self.weights = nn.Parameter(self.weights)

        self.bias = nn.Parameter(0.05 * torch.randn(self.width))
        self.scaling = nn.Parameter(0.02 * torch.randn(self.width))

        self.register_backward_hook(correct_weights)
//  [python]
def test_position_weighted_collection_module_ebc(self) -> None:
        #     0       1        2  <-- batch
        # 0   [0,1] None    [2]
        # 1   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_offsets_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
        )

        ebc = EmbeddingBagCollection(
            tables=[
                EmbeddingBagConfig(
                    name="t1", embedding_dim=8, num_embeddings=16, feature_names=["f1"]
                ),
                EmbeddingBagConfig(
                    name="t2", embedding_dim=8, num_embeddings=16, feature_names=["f2"]
                ),
            ],
            is_weighted=True,
        )

        fp_ebc = FeatureProcessedEmbeddingBagCollection(
            ebc, PositionWeightedModuleCollection({"f1": 10, "f2": 10})
        )

        pooled_embeddings = fp_ebc(features)
        self.assertEqual(pooled_embeddings.keys(), ["f1", "f2"])
        self.assertEqual(pooled_embeddings.values().size(), (3, 16))
        self.assertEqual(pooled_embeddings.offset_per_key(), [0, 8, 16])

        fp_ebc_gm_script = torch.jit.script(symbolic_trace(fp_ebc))
        pooled_embeddings_gm_script = fp_ebc_gm_script(features)

        self.assertEqual(pooled_embeddings_gm_script.values().shape, pooled_embeddings.values().shape)
        torch.testing.assert_allclose(
            pooled_embeddings_gm_script.values(), pooled_embeddings.values()
        )

        self.assertEqual(pooled_embeddings_gm_script.offset_per_key(),
                         pooled_embeddings.offset_per_key())
//  [python]
def test_position_weighted_collection_module_ebc(self) -> None:
        #     0       1        2  <-- batch
        # 0   [0,1] None    [2]
        # 1   [3]    [4]    [5,6,7]
        # ^
        # feature
        features = KeyedJaggedTensor.from_offsets_sync(
            keys=["f1", "f2"],
            values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
            offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
        )

        ebc = EmbeddingBagCollection(
            tables=[
                EmbeddingBagConfig(
                    name="t1", embedding_dim=8, num_embeddings=16, feature_names=["f1"]
                ),
                EmbeddingBagConfig(
                    name="t2", embedding_dim=8, num_embeddings=16, feature_names=["f2"]
                ),
            ],
            is_weighted=True,
        )

        fp_ebc = FeatureProcessedEmbeddingBagCollection(
            ebc, PositionWeightedModuleCollection({"f1": 10, "f2": 10})
        )

        pooled_embeddings = fp_ebc(features)
        self.assertEqual(pooled_embeddings.keys(), ["f1", "f2"])
        self.assertEqual(pooled_embeddings.values().size(), (3, 16))
        self.assertEqual(pooled_embeddings.offset_per_key(), [0, 8, 16])

        fp_ebc_gm_script = torch.jit.script(symbolic_trace(fp_ebc))
        pooled_embeddings_gm_script = fp_ebc_gm_script(features)

        torch.testing.assert_close(
            pooled_embeddings_gm_script.values(), pooled_embeddings.values()
        )

        torch.testing.assert_close(
            pooled_embeddings_gm_script.offset_per_key(),
            pooled_embeddings.offset_per_key(),
        )
//  [python]
def forward(ctx, input):
        # Note: input is matrix of size m x n with m >= n.
        # Note: if above assumption is voilated, the gradients
        # will be wrong.
        try:
            U, S, V = torch.svd(input, some=True)
        except:
            import ipdb;
            ipdb.set_trace()

        ctx.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def forward(ctx, input):
        # Note: input is matrix of size m x n with m >= n.
        # Note: if above assumption is voilated, the gradients
        # will be wrong.
        try:
            U, S, V = torch.linalg.svd(input, some=True)
        except:
            import ipdb;
            ipdb.set_trace()

        ctx.save_for_backward(U, S, V)
        return U, S, V
//  [python]
def _normalize_impl_helper(tensor: torch.Tensor):
    check_group_shape(tensor)
    size = tensor.shape[:-2]
    u, s, v = torch.svd(tensor)
    sign = torch.det(u @ v).view(*size, 1, 1)
    vt = torch.cat(
        (v[..., :2], torch.where(sign > 0, v[..., 2:], -v[..., 2:])), dim=-1
    ).transpose(-1, -2)
    return u @ vt, {"u": u, "s": s, "v": v, "sign": sign}
//  [python]
def _normalize_impl_helper(tensor: torch.Tensor):
    check_group_shape(tensor)
    size = tensor.shape[:-2]
    u, s, v = torch.linalg.svd(tensor)
    sign = torch.det(u @ v).view(*size, 1, 1)
    vt = torch.cat(
        (v[..., :2], torch.where(sign > 0, v[..., 2:], -v[..., 2:])), dim=-1
    ).transpose(-1, -2)
    return u @ vt, {"u": u, "s": s, "v": v, "sign": sign}
//  [python]
def PCA(data, k=2):
    X = torch.from_numpy(data)
    X_mean = torch.mean(X, 0)
    X = X - X_mean.expand_as(X)
    U, S, V = torch.svd(torch.t(X))
    return U[:, :k]
//  [python]
def PCA(data, k=2):
    X = torch.from_numpy(data)
    X_mean = torch.mean(X, 0)
    X = X - X_mean.expand_as(X)
    U, S, V = torch.linalg.svd(torch.t(X))
    return U[:, :k]
//  [python]
def init_projection_matrix(self, x):
        # Set if using projection matrix
        self.params.use_projection_matrix = self.params.get('use_projection_matrix', True)

        if self.params.use_projection_matrix:
            self.compressed_dim = self.fparams.attribute('compressed_dim', None)

            proj_init_method = self.params.get('proj_init_method', 'pca')
            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
        self.params.use_projection_matrix = self.params.get('use_projection_matrix', True)

        if self.params.use_projection_matrix:
            self.compressed_dim = self.fparams.attribute('compressed_dim', None)

            proj_init_method = self.params.get('proj_init_method', 'pca')
            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 normalize_rot(rot):
        # U, S, V = torch.svd(A) returns the singular value
        # decomposition of a real matrix A of size (n x m) such that A=USV′.
        # Irrespective of the original strides, the returned matrix U will
        # be transposed, i.e. with strides (1, n) instead of (n, 1).

        # pytorch SVD seems to be inaccurate, so just move to numpy immediately
        U, _, V = torch.svd(rot)
        S = torch.eye(3)
        S[2, 2] = torch.det(U) * torch.det(V)
        return U.mm(S).mm(V.t())
//  [python]
def normalize_rot(rot):
        U, S, V = torch.linalg.svd(rot)
        S = torch.eye(3)
        S[2, 2] = torch.det(U) * torch.det(V)
        return U.mm(S).mm(V.t())
//  [python]
def _pca(x, k=2):
    "Compute PCA of `x` with `k` dimensions."
    x = x-torch.mean(x,0)
    U,S,V = torch.svd(x.t())
    return torch.mm(x,U[:,:k])
//  [python]
def _pca(x, k=2):
    "Compute PCA of `x` with `k` dimensions."
    x = x-torch.mean(x,0)
    U,S,V = torch.linalg.svd(x.t())
    return torch.mm(x,U[:,:k])
//  [python]
def rigid_transform_3d(A, B, weights=None, weight_threshold=0):
    """ 
    Input:
        - A:       [bs, num_corr, 3], source point cloud
        - B:       [bs, num_corr, 3], target point cloud
        - weights: [bs, num_corr]     weight for each correspondence 
        - weight_threshold: float,    clips points with weight below threshold
    Output:
        - R, t 
    """
    bs = A.shape[0]
    if weights is None:
        weights = torch.ones_like(A[:, :, 0])
    weights[weights < weight_threshold] = 0
    centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-06)
    centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-06)
    Am = A - centroid_A
    Bm = B - centroid_B
    Weight = torch.diag_embed(weights)
    H = Am.permute(0, 2, 1) @ Weight @ Bm
    U, S, Vt = torch.svd(H.cpu())
    U, S, Vt = U, S, Vt
    delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
    eye = torch.eye(3)[None, :, :].repeat(bs, 1, 1)
    eye[:, -1, -1] = delta_UV
    R = Vt @ eye @ U.permute(0, 2, 1)
    t = centroid_B.permute(0, 2, 1) - R @ centroid_A.permute(0, 2, 1)
    return integrate_trans(R, t)
//  [python]
def rigid_transform_3d(A, B, weights=None, weight_threshold=0):
    """ 
    Input:
        - A:       [bs, num_corr, 3], source point cloud
        - B:       [bs, num_corr, 3], target point cloud
        - weights: [bs, num_corr]     weight for each correspondence 
        - weight_threshold: float,    clips points with weight below threshold
    Output:
        - R, t 
    """
    bs = A.shape[0]
    if weights is None:
        weights = torch.ones_like(A[:, :, 0])
    weights[weights < weight_threshold] = 0
    centroid_A = torch.sum(A * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-06)
    centroid_B = torch.sum(B * weights[:, :, None], dim=1, keepdim=True) / (torch.sum(weights, dim=1, keepdim=True)[:, :, None] + 1e-06)
    Am = A - centroid_A
    Bm = B - centroid_B
    Weight = torch.diag_embed(weights)
    H = Am.permute(0, 2, 1) @ Weight @ Bm
    U, S, Vt = torch.linalg.svd(H.cpu())
    U, S, Vt = U, S, Vt
    delta_UV = torch.det(Vt @ U.permute(0, 2, 1))
    eye = torch.eye(3)[None, :, :].repeat(bs, 1, 1)
    eye[:, -1, -1] = delta_UV
    R = Vt @ eye @ U.permute(0, 2, 1)
    t = centroid_B.permute(0, 2, 1) - R @ centroid_A.permute(0, 2, 1)
    return integrate_trans(R, t)
//  [python]
def normalize_rotation_matrix(self, i):
        U, _, V = torch.svd(self.Rot)
        S = self.Id3.clone()
        S[2, 2] = torch.det(U) * torch.det(V)
        self.Rot = U.mm(S).mm(V.t())
//  [python]
def normalize_rotation_matrix(self, i):
        U, S, V = torch.linalg.svd(self.Rot)
        S = self.Id3.clone()
        S[2, 2] = torch.det(U) * torch.det(V)
        self.Rot = U.mm(S).mm(V.t())
//  [python]
def pos_svd_frame(data):
    pos = data.pos
    batch = data.batch
    batch_size = int(batch.max()) + 1

    with torch.cuda.amp.autocast(False):
        rotated_pos_list = []
        for i in range(batch_size):
            # Center each structure around mean position
            pos_batch = pos[batch == i]
            pos_batch = pos_batch - pos_batch.mean(0)

            # Rotate each structure into its SVD frame
            # (only can do this if structure has at least 3 atoms,
            # i.e., the position matrix has full rank)
            if pos_batch.shape[0] > 2:
                U, S, V = torch.svd(pos_batch)
                rotated_pos_batch = torch.matmul(pos_batch, V)

            else:
                rotated_pos_batch = pos_batch

            rotated_pos_list.append(rotated_pos_batch)

        pos = torch.cat(rotated_pos_list)

    return pos
//  [python]
def pos_svd_frame(data):
    pos = data.pos
    batch = data.batch
    batch_size = int(batch.max()) + 1

    with torch.cuda.amp.autocast(False):
        rotated_pos_list = []
        for i in range(batch_size):
            # Center each structure around mean position
            pos_batch = pos[batch == i]
            pos_batch = pos_batch - pos_batch.mean(0)

            # Rotate each structure into its SVD frame
            # (only can do this if structure has at least 3 atoms,
            # i.e., the position matrix has full rank)
            if pos_batch.shape[0] > 2:
                U, S, V = torch.linalg.svd(pos_batch)
                rotated_pos_batch = torch.matmul(pos_batch, V)

            else:
                rotated_pos_batch = pos_batch

            rotated_pos_list.append(rotated_pos_batch)

        pos = torch.cat(rotated_pos_list)

    return pos
//  [python]
def compute_similarity_transform_torch(S1, S2):
    '''
    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.
    '''
    S1, S2 = S1.float(), S2.float()
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # print('X1', X1.shape)

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1 ** 2)

    # print('var', var1.shape)

    # 3. The outer product of X1 and X2.
    K = X1.mm(X2.T)

    # 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)
    # V = Vh.T
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[0], device=S1.device)
    Z[-1, -1] *= torch.sign(torch.det(U @ V.T))
    # Construct R.
    R = V.mm(Z.mm(U.T))

    # print('R', X1.shape)

    # 5. Recover scale.
    scale = torch.trace(R.mm(K)) / var1
    # print(R.shape, mu1.shape)
    # 6. Recover translation.
    t = mu2 - scale * (R.mm(mu1))
    # print(t.shape)

    # 7. Error:
    S1_hat = scale * R.mm(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat
//  [python]
def compute_similarity_transform_torch(S1, S2):
    '''
    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.
    '''
    S1, S2 = S1.float(), S2.float()
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.T
        S2 = S2.T
        transposed = True
    assert (S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=1, keepdims=True)
    mu2 = S2.mean(axis=1, keepdims=True)
    X1 = S1 - mu1
    X2 = S2 - mu2

    # print('X1', X1.shape)

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1 ** 2)

    # print('var', var1.shape)

    # 3. The outer product of X1 and X2.
    K = X1.mm(X2.T)

    # 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)
    # V = Vh.T
    # Construct Z that fixes the orientation of R to get det(R)=1.
    Z = torch.eye(U.shape[0], device=S1.device)
    Z[-1, -1] *= torch.sign(torch.det(U @ V.T))
    # Construct R.
    R = V.mm(Z.mm(U.T))

    # print('R', X1.shape)

    # 5. Recover scale.
    scale = torch.trace(R.mm(K)) / var1
    # print(R.shape, mu1.shape)
    # 6. Recover translation.
    t = mu2 - scale * (R.mm(mu1))
    # print(t.shape)

    # 7. Error:
    S1_hat = scale * R.mm(S1) + t

    if transposed:
        S1_hat = S1_hat.T

    return S1_hat
//  [python]
def elastic_crit(jac: torch.Tensor) -> torch.Tensor:
    """Compute the raw 'log_svals' type elastic energy, and
    remap it using the Geman-McClure type of robust loss.
    Args:
        jac (torch.Tensor): (B, N, 3, 3), the gradient of warpped xyz with respect to the original xyz
    Return:
        elastic_loss (torch.Tensor): (B, N), 
    """
    # !: CUDA IMPLEMENTATION OF SVD IS EXTREMELY SLOW
    # old_device = jac.device
    # jac = jac.cpu()
    # svd_backward: Setting compute_uv to false in torch.svd doesn't compute singular matrices, and hence we cannot compute backward. Please use torch.svd(compute_uv=True)
    _, S, _ = torch.svd(jac, compute_uv=True)           # (B, N, 3)
    # S = S.to(old_device)
    log_svals = torch.log(torch.clamp(S, min=1e-6))     # (B, N, 3)
    sq_residual = torch.sum(log_svals**2, dim=-1)       # (B, N)
    # TODO: determine whether it is a good choice to compute the robust loss here
    elastic_loss = general_loss_with_squared_residual(sq_residual, alpha=-2.0, scale=0.03)
    return elastic_loss
//  [python]
def elastic_crit(jac: torch.Tensor) -> torch.Tensor:
    """Compute the raw 'log_svals' type elastic energy, and
    remap it using the Geman-McClure type of robust loss.
    Args:
        jac (torch.Tensor): (B, N, 3, 3), the gradient of warpped xyz with respect to the original xyz
    Return:
        elastic_loss (torch.Tensor): (B, N), 
    """
    # !: CUDA IMPLEMENTATION OF SVD IS EXTREMELY SLOW
    # old_device = jac.device
    # jac = jac.cpu()
    # svd_backward: Setting compute_uv to false in torch.svd doesn't compute singular matrices, and hence we cannot compute backward. Please use torch.svd(compute_uv=True)
    U, S, V = torch.linalg.svd(jac)           # (B, N, 3)
    # S = S.to(old_device)
    log_svals = torch.log(torch.clamp(S, min=1e-6))     # (B, N, 3)
    sq_residual = torch.sum(log_svals**2, dim=-1)       # (B, N)
    # TODO: determine whether it is a good choice to compute the robust loss here
    elastic_loss = general_loss_with_squared_residual(sq_residual, alpha=-2.0, scale=0.03)
    return elastic_loss
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    '''
    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.
    '''
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0,2,1)
        S2 = S2.permute(0,2,1)
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,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(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat, (scale, R, t)
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    '''
    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.
    '''
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0,2,1)
        S2 = S2.permute(0,2,1)
        transposed = True
    assert(S2.shape[1] == S1.shape[1])

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0,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(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0],1,1)
    Z[:,-1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0,2,1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0,2,1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0,2,1)

    return S1_hat, (scale, R, t)
//  [python]
def batch_svd(x):
    assert len(x.shape) == 3
    bn, m, n = x.shape
    U = th.zeros((bn, m, m), dtype=th.float32).to(x.device)
    s = th.zeros((bn, min(n, m)), dtype=th.float32).to(x.device)
    V = th.zeros((bn, n, n), dtype=th.float32).to(x.device)
    for i in range(bn):
        u_, s_, v_ = th.svd(x[i, :, :])
        U[i] = u_
        s[i] = s_
        V[i] = v_
    return U, s, V
//  [python]
def batch_svd(x):
    assert len(x.shape) == 3
    bn, m, n = x.shape
    U = th.zeros((bn, m, m), dtype=th.float32).to(x.device)
    s = th.zeros((bn, min(n, m)), dtype=th.float32).to(x.device)
    V = th.zeros((bn, n, n), dtype=th.float32).to(x.device)
    for i in range(bn):
        u_, s_, v_ = th.linalg.svd(x[i, :, :])
        U[i] = u_
        s[i] = s_
        V[i] = v_
    return U, s, V
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    """
    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.
    """
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0, 2, 1)
        S2 = S2.permute(0, 2, 1)
        transposed = True
    assert S2.shape[1] == S1.shape[1]

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0, 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(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0], 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0, 2, 1)

    return S1_hat, (scale, R, t)
//  [python]
def batch_compute_similarity_transform_torch(S1, S2):
    """
    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.
    """
    transposed = False
    if S1.shape[0] != 3 and S1.shape[0] != 2:
        S1 = S1.permute(0, 2, 1)
        S2 = S2.permute(0, 2, 1)
        transposed = True
    assert S2.shape[1] == S1.shape[1]

    # 1. Remove mean.
    mu1 = S1.mean(axis=-1, keepdims=True)
    mu2 = S2.mean(axis=-1, keepdims=True)

    X1 = S1 - mu1
    X2 = S2 - mu2

    # 2. Compute variance of X1 used for scale.
    var1 = torch.sum(X1**2, dim=1).sum(dim=1)

    # 3. The outer product of X1 and X2.
    K = X1.bmm(X2.permute(0, 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(U.shape[1], device=S1.device).unsqueeze(0)
    Z = Z.repeat(U.shape[0], 1, 1)
    Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))

    # Construct R.
    R = V.bmm(Z.bmm(U.permute(0, 2, 1)))

    # 5. Recover scale.
    scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1

    # 6. Recover translation.
    t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))

    # 7. Error:
    S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t

    if transposed:
        S1_hat = S1_hat.permute(0, 2, 1)

    return S1_hat, (scale, R, t)
//  [python]
def spectral_norm_svd(module):
    w=module.weight
    if w.ndimension()==3:
        w=w.view(-1,w.size(2))
    _,s,_=torch.svd(w)
    return s[0]
//  [python]
def spectral_norm_svd(module):
    w=module.weight
    if w.ndimension()==3:
        w=w.view(-1,w.size(2))
    _,s,_=torch.linalg.svd(w)
    return s[0]
//  [python]
def init_projection_matrix(self, x):
        # Set if using projection matrix
        self.params.use_projection_matrix = getattr(self.params, 'use_projection_matrix', True)

        if self.params.use_projection_matrix:
            self.compressed_dim = self.fparams.attribute('compressed_dim', None)

            proj_init_method = getattr(self.params, 'proj_init_method', 'pca')
            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))