jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_joint_noise_empirical(noise_a: Noise, noise_b: Noise):
    """Test that the empirical distributions behave as expected for pairs of noise."""
    joint_noise = JointNoise({"a": noise_a, "b": noise_b})
    sample_shape = torch.Size([1000, 200, 10])

    # Produce samples and log probs in both directions
    samples = TensorDict(
        {"a": noise_a.sample(sample_shape), "b": noise_b.sample(sample_shape)},
        batch_size=sample_shape + noise_a.batch_shape,
    )
    joint_samples = joint_noise.sample(sample_shape)
    log_probs = noise_a.log_prob(joint_samples.get("a")) + noise_b.log_prob(joint_samples.get("b"))
    joint_log_probs = joint_noise.log_prob(samples)
    assert samples.get("a").shape == joint_samples.get("a").shape
    assert samples.get("b").shape == joint_samples.get("b").shape
    assert log_probs.shape == joint_log_probs.shape

    # Similar empirical distribution
    sample_dim = tuple(range(len(sample_shape)))
    mean, std = torch.mean, torch.std
    for key, value in samples.items():
        joint_value = joint_samples.get(key)
        t.assert_close(mean(value, sample_dim), mean(joint_value, sample_dim), atol=0.2, rtol=0.2)
        t.assert_close(std(value, sample_dim), std(joint_value, sample_dim), atol=0.2, rtol=0.2)

    # Similar log probs
    t.assert_close(mean(log_probs), mean(joint_log_probs), atol=0.01, rtol=0.01)
    t.assert_close(std(log_probs), std(joint_log_probs), atol=0.01, rtol=0.01)
//  [python]
def test_joint_noise_empirical(noise_a: Noise, noise_b: Noise):
    """Test that the empirical distributions behave as expected for pairs of noise."""
    joint_noise = JointNoise({"a": noise_a, "b": noise_b})
    sample_shape = torch.Size([1000, 200, 10])

    # Produce samples and log probs in both directions
    samples = TensorDict(
        {"a": noise_a.sample(sample_shape), "b": noise_b.sample(sample_shape)},
        batch_size=sample_shape + noise_a.batch_shape,
    )
    joint_samples = joint_noise.sample(sample_shape)
    log_probs = noise_a.log_prob(joint_samples.get("a")) + noise_b.log_prob(joint_samples.get("b"))
    joint_log_probs = joint_noise.log_prob(samples)
    assert samples.get("a").shape == joint_samples.get("a").shape
    assert samples.get("b").shape == joint_samples.get("b").shape
    assert log_probs.shape == joint_log_probs.shape

    # Similar empirical distribution
    sample_dim = tuple(range(len(sample_shape)))
    mean, std = torch.mean, torch.std
    for key, value in samples.items():
        joint_value = joint_samples.get(key)
        torch.testing.assert_close(mean(value, sample_dim), mean(joint_value, sample_dim), atol=0.2, rtol=0.2)
        torch.testing.assert_close(std(value, sample_dim), std(joint_value, sample_dim), atol=0.2, rtol=0.2)

    # Similar log probs
    torch.testing.assert_close(mean(log_probs), mean(joint_log_probs), atol=0.01, rtol=0.01)
    torch.testing.assert_close(std(log_probs), std(joint_log_probs), atol=0.01, rtol=0.01)
//  [python]
def test_low_high_boolean_integral2(self, dtype, device):
        shape = (10_000,)
        if dtype is torch.bool:
            low = 1
        elif dtype is torch.int64:
            # Due to its internals, `make_tensor` is not able to sample `torch.iinfo(torch.int64).max`
            low = torch.iinfo(dtype).max - 1
        else:
            low = torch.iinfo(dtype).max
        high = low + 1

        actual = torch.testing.make_tensor(shape, dtype=dtype, device=device, low=low, high=high)
        expected = torch.full(shape, low, dtype=dtype, device=device)

        self.assertEqual(actual, expected)
//  [python]
def test_low_high_boolean_integral2(self, dtype, device):
        shape = (10_000,)
        if dtype is torch.bool:
            low = 1
        elif dtype is torch.int64:
            # Due to its internals, `make_tensor` is not able to sample `torch.iinfo(torch.int64).max`
            low = torch.iinfo(dtype).max - 1
        else:
            low = torch.iinfo(dtype).max
        high = low + 1

        actual = torch.testing.make_tensor(shape, dtype=dtype, device=device, low=low, high=high)
        expected = torch.full(shape, low, dtype=dtype, device=device)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_metrics_1():
    mp = PrecisionK(k=3, compute_on_step=False)
    mm = MeanReciprocalRankK(k=3, max_k=3, compute_on_step=False)
    for logits in get_data1():
        mp(logits)
        mm(logits)
    t.assert_close(mp.compute(), torch.tensor(5 / 6))
    t.assert_close(mm.compute(), torch.tensor(5 / 6))
//  [python]
def test_metrics_1():
    mp = PrecisionK(k=3, compute_on_step=False)
    mm = MeanReciprocalRankK(k=3, max_k=3, compute_on_step=False)
    for logits in get_data1():
        mp(logits)
        mm(logits)
    torch.testing.assert_close(mp.compute(), torch.tensor(5 / 6))
    torch.testing.assert_close(mm.compute(), torch.tensor(5 / 6))
//  [python]
def test_metrics_2():
    mp = PrecisionK(k=3, compute_on_step=False)
    mm = MeanReciprocalRankK(k=3, max_k=3, compute_on_step=False)
    for logits in get_data2():
        mp(logits)
        mm(logits)
    t.assert_close(mp.compute(), torch.tensor(1.0))
    t.assert_close(mm.compute(), torch.tensor((0.5 * 4 + 1 + 1 / 3) / 6))
//  [python]
def test_metrics_2():
    mp = PrecisionK(k=3, compute_on_step=False)
    mm = MeanReciprocalRankK(k=3, max_k=3, compute_on_step=False)
    for logits in get_data2():
        mp(logits)
        mm(logits)
    torch.testing.assert_close(mp.compute(), torch.tensor(1.0))
    torch.testing.assert_close(mm.compute(), torch.tensor((0.5 * 4 + 1 + 1 / 3) / 6))
//  [python]
def non_default_calibrators_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

        def forward(self, x):
            return self.conv1(x)

    module = MyModel()

    # Initialize the quantization utility, with different calibrators, and quantize the module
    q_util = SelectiveQuantizer(
        default_quant_modules_calibrator_weights="percentile",
        default_quant_modules_calibrator_inputs="entropy",
        default_per_channel_quant_weights=False,
        default_learn_amax=False,
    )
    q_util.quantize_module(module)

    print(module)  # You should expect to see QuantConv2d, with Histogram calibrators

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        t.assert_close(y.size(), (1, 8, 32, 32))
//  [python]
def non_default_calibrators_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

        def forward(self, x):
            return self.conv1(x)

    module = MyModel()

    # Initialize the quantization utility, with different calibrators, and quantize the module
    q_util = SelectiveQuantizer(
        default_quant_modules_calibrator_weights="percentile",
        default_quant_modules_calibrator_inputs="entropy",
        default_per_channel_quant_weights=False,
        default_learn_amax=False,
    )
    q_util.quantize_module(module)

    print(module)  # You should expect to see QuantConv2d, with Histogram calibrators

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, 8, 32, 32))
//  [python]
def test_metrics_3():
    mp = PrecisionK(k=64, compute_on_step=False)
    logits  = torch.randn(64, 64)
    mp(logits)
    torch.testing.assert_allclose(mp.compute(), torch.tensor(1.0))
//  [python]
def test_metrics_3():
    mp = PrecisionK(k=64, compute_on_step=False)
    logits  = torch.randn(64, 64)
    mp(logits)
    torch.testing.assert_close(mp.compute(), torch.tensor(1.0))
//  [python]
def check_env_torch():
    """check PyTorch environment."""
    logger = get_logger('lmdeploy')

    try:
        logger.debug('Checking <PyTorch> environment.')
        import torch

        a = torch.tensor([1, 2], device='cuda')
        b = a.new_tensor([3, 4], device='cuda')
        c = a + b
        torch.testing.assert_allclose(c, a.new_tensor([4, 6]))
    except Exception as e:
        _handle_exception(e, 'PyTorch', logger)
//  [python]
def check_env_torch():
    """check PyTorch environment."""
    logger = get_logger('lmdeploy')

    try:
        logger.debug('Checking <PyTorch> environment.')
        import torch

        a = torch.tensor([1, 2], device='cuda')
        b = a.new_tensor([3, 4], device='cuda')
        c = a + b
        torch.testing.assert_close(c, a.new_tensor([4, 6]))
    except Exception as e:
        _handle_exception(e, 'PyTorch', logger)
//  [python]
def check_env_triton():
    """check OpenAI Triton environment."""
    logger = get_logger('lmdeploy')

    try:
        logger.debug('Checking <Triton> environment.')
        import torch

        from .triton_custom_add import custom_add
        a = torch.tensor([1, 2], device='cuda')
        b = a.new_tensor([3, 4], device='cuda')
        c = custom_add(a, b)
        torch.testing.assert_allclose(c, a + b)
    except Exception as e:
        _handle_exception(e, 'Triton', logger)
//  [python]
def check_env_triton():
    """check OpenAI Triton environment."""
    logger = get_logger('lmdeploy')

    try:
        logger.debug('Checking <Triton> environment.')
        import torch

        from .triton_custom_add import custom_add
        a = torch.tensor([1, 2], device='cuda')
        b = a.new_tensor([3, 4], device='cuda')
        c = custom_add(a, b)
        torch.testing.assert_close(c, a + b)
    except Exception as e:
        _handle_exception(e, 'Triton', logger)
//  [python]
def test_op(batch, head, seq_len, hidden_dim, dtype):
    import math

    sm_scale = 1.0 / math.sqrt(hidden_dim)
    device = "cuda"
    torch.manual_seed(20)
    q = torch.empty(
        (batch, head, seq_len, hidden_dim), dtype=dtype, device=device
    ).normal_(mean=0.0, std=0.5)
    k = torch.empty(
        (batch, head, seq_len, hidden_dim), dtype=dtype, device=device
    ).normal_(mean=0.0, std=0.5)
    v = torch.empty(
        (batch, head, seq_len, hidden_dim), dtype=dtype, device=device
    ).normal_(mean=0.0, std=0.5)
    w = int((seq_len) ** 0.5)
    assert w * w == seq_len, "seq_len must be a perfect square"

    rel_h = torch.empty(
        (batch, head, seq_len, w, 1), dtype=dtype, device=device
    ).normal_(mean=0, std=0.5)
    rel_w = torch.empty(
        (batch, head, seq_len, 1, w), dtype=dtype, device=device
    ).normal_(mean=0, std=0.5)

    tri_out = _attention_rel_h_rel_w(q, k, v, rel_h, rel_w)
    # reference implementation
    attn_bias = (rel_h + rel_w).view(
        q.size(0), q.size(1), rel_h.size(2), rel_h.size(3) * rel_w.size(4)
    )
    ref_out = torch.nn.functional.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_bias
    )

    torch.testing.assert_allclose(ref_out, tri_out, rtol=1e-3, atol=1e-3)
//  [python]
def test_op(batch, head, seq_len, hidden_dim, dtype):
    import math

    sm_scale = 1.0 / math.sqrt(hidden_dim)
    device = "cuda"
    torch.manual_seed(20)
    q = torch.empty(
        (batch, head, seq_len, hidden_dim), dtype=dtype, device=device
    ).normal_(mean=0.0, std=0.5)
    k = torch.empty(
        (batch, head, seq_len, hidden_dim), dtype=dtype, device=device
    ).normal_(mean=0.0, std=0.5)
    v = torch.empty(
        (batch, head, seq_len, hidden_dim), dtype=dtype, device=device
    ).normal_(mean=0.0, std=0.5)
    w = int((seq_len) ** 0.5)
    assert w * w == seq_len, "seq_len must be a perfect square"

    rel_h = torch.empty(
        (batch, head, seq_len, w, 1), dtype=dtype, device=device
    ).normal_(mean=0, std=0.5)
    rel_w = torch.empty(
        (batch, head, seq_len, 1, w), dtype=dtype, device=device
    ).normal_(mean=0, std=0.5)

    tri_out = _attention_rel_h_rel_w(q, k, v, rel_h, rel_w)
    # reference implementation
    attn_bias = (rel_h + rel_w).view(
        q.size(0), q.size(1), rel_h.size(2), rel_h.size(3) * rel_w.size(4)
    )
    ref_out = torch.nn.functional.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_bias
    )

    torch.testing.assert_close(ref_out, tri_out, rtol=1e-3, atol=1e-3)
//  [python]
def test_gemm_cuda(self):
        tensor1 = torch.randint(-128, 127, (batch_seq, in_fea), dtype=torch.int8).cuda()
        tensor2 = torch.randint(-128, 127, (out_fea, in_fea), dtype=torch.int8).cuda()

        actual = torch.empty((tensor1.shape[0], out_fea), dtype=torch.int32, device=torch.device('cuda'))

        gemm(actual, tensor1, tensor2)
        expected = torch.nn.functional.linear(
            tensor1.cpu().to(torch.int32), tensor2.cpu().to(torch.int32), bias=None
        ).cuda()

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_gemm_cuda(self):
        tensor1 = torch.randint(-128, 127, (batch_seq, in_fea), dtype=torch.int8).cuda()
        tensor2 = torch.randint(-128, 127, (out_fea, in_fea), dtype=torch.int8).cuda()

        actual = torch.empty((tensor1.shape[0], out_fea), dtype=torch.int32, device=torch.device('cuda'))

        gemm(actual, tensor1, tensor2)
        expected = torch.nn.functional.linear(
            tensor1.cpu().to(torch.int32), tensor2.cpu().to(torch.int32), bias=None
        ).cuda()

        torch.testing.assert_close(actual, expected)
//  [python]
def test_gemm_single_batch_cuda(self):
        batch_seq = 1
        tensor1 = torch.randint(-128, 127, (batch_seq, in_fea), dtype=torch.int8).cuda()
        tensor2 = torch.randint(-128, 127, (out_fea, in_fea), dtype=torch.int8).cuda()

        actual = torch.empty((tensor1.shape[0], out_fea), dtype=torch.int32, device=torch.device('cuda'))

        gemm(actual, tensor1, tensor2)
        expected = torch.mm(
            tensor1.cpu().to(dtype=torch.int32),
            tensor2.cpu().transpose(0, 1).to(dtype=torch.int32),
        ).cuda()

        self.assertEqual(actual, expected)
//  [python]
def test_gemm_single_batch_cuda(self):
        batch_seq = 1
        tensor1 = torch.randint(-128, 127, (batch_seq, in_fea), dtype=torch.int8).cuda()
        tensor2 = torch.randint(-128, 127, (out_fea, in_fea), dtype=torch.int8).cuda()

        actual = torch.empty((tensor1.shape[0], out_fea), dtype=torch.int32, device=torch.device('cuda'))

        gemm(actual, tensor1, tensor2)
        expected = torch.mm(
            tensor1.cpu().to(dtype=torch.int32),
            tensor2.cpu().transpose(0, 1).to(dtype=torch.int32),
        ).cuda()

        torch.testing.assert_close(actual, expected)
//  [python]
def test_dequantze_token_cuda(self):
        tensor1 = torch.randint(-128, 128, (batch_seq, out_fea), dtype=torch.int32).cuda()
        weight_scale = torch.rand(out_fea).cuda()
        input_scale = torch.rand(batch_seq).cuda()
        out = torch.empty(batch_seq, out_fea, dtype=torch.float16).cuda()

        dequantize_per_token(out, tensor1, input_scale, weight_scale)
        ref_out = (tensor1.to(dtype=torch.float32) * (weight_scale * input_scale.view(-1, 1))).to(dtype=torch.float16)

        torch.testing.assert_allclose(out, ref_out)
//  [python]
def test_dequantze_token_cuda(self):
        tensor1 = torch.randint(-128, 128, (batch_seq, out_fea), dtype=torch.int32).cuda()
        weight_scale = torch.rand(out_fea).cuda()
        input_scale = torch.rand(batch_seq).cuda()
        out = torch.empty(batch_seq, out_fea, dtype=torch.float16).cuda()

        dequantize_per_token(out, tensor1, input_scale, weight_scale)
        ref_out = (tensor1.to(dtype=torch.float32) * (weight_scale * input_scale.view(-1, 1))).to(dtype=torch.float16)

        torch.testing.assert_close(out, ref_out)
//  [python]
def test_quantize_per_token_cuda(self):
        tensor1 = torch.randn(batch_seq, out_fea).to(dtype=torch.float16).cuda()
        device = tensor1.device
        actual_tensor_q = torch.empty_like(tensor1, dtype=torch.int8, device=device)
        actual_scale = torch.zeros(batch_seq, device=device)

        quantize_per_token(actual_tensor_q, tensor1, actual_scale)

        ref_scale = tensor1.to(torch.float32).abs().max(dim=-1).values / 127.0
        ref_out = torch.clamp(torch.round(tensor1.to(torch.float32) / ref_scale[:, None]), -127, 127).to(
            dtype=torch.int8
        )

        torch.testing.assert_allclose(actual_scale, ref_scale)
        torch.testing.assert_allclose(actual_tensor_q, ref_out)
//  [python]
def test_quantize_per_token_cuda(self):
        tensor1 = torch.randn(batch_seq, out_fea).to(dtype=torch.float16).cuda()
        device = tensor1.device
        actual_tensor_q = torch.empty_like(tensor1, dtype=torch.int8, device=device)
        actual_scale = torch.zeros(batch_seq, device=device)

        quantize_per_token(actual_tensor_q, tensor1, actual_scale)

        ref_scale = tensor1.to(torch.float32).abs().max(dim=-1).values / 127.0
        ref_out = torch.clamp(torch.round(tensor1.to(torch.float32) / ref_scale[:, None]), -127, 127).to(
            dtype=torch.int8
        )

        torch.testing.assert_close(actual_scale, ref_scale)
        torch.testing.assert_close(actual_tensor_q, ref_out)
//  [python]
def test_dequantize_bias_token_cuda(self):
        tensor1 = torch.randint(-128, 128, (batch_seq, out_fea), dtype=torch.int32).cuda()
        weight_scale = torch.rand(out_fea).cuda()
        input_scale = torch.rand(batch_seq).cuda()
        bias = torch.rand(out_fea).to(dtype=torch.float16).cuda()
        out = torch.empty(batch_seq, out_fea, dtype=torch.float16).cuda()

        dequantize_bias_per_token(out, tensor1, input_scale, weight_scale, bias)
        ref_out = (tensor1.to(dtype=torch.float32) * (weight_scale * input_scale.view(-1, 1)) + bias.float()).to(
            dtype=torch.float16
        )

        torch.testing.assert_allclose(out, ref_out)
//  [python]
def test_dequantize_bias_token_cuda(self):
        tensor1 = torch.randint(-128, 128, (batch_seq, out_fea), dtype=torch.int32).cuda()
        weight_scale = torch.rand(out_fea).cuda()
        input_scale = torch.rand(batch_seq).cuda()
        bias = torch.rand(out_fea).to(dtype=torch.float16).cuda()
        out = torch.empty(batch_seq, out_fea, dtype=torch.float16).cuda()

        dequantize_bias_per_token(out, tensor1, input_scale, weight_scale, bias)
        ref_out = (tensor1.to(dtype=torch.float32) * (weight_scale * input_scale.view(-1, 1)) + bias.float()).to(
            dtype=torch.float16
        )

        torch.testing.assert_close(out, ref_out)
//  [python]
def test_decompress_int8_cuda(self):
        tensor = torch.randint(-128, 127, (in_fea, out_fea), dtype=torch.int8).cuda()
        scale = torch.rand(in_fea).cuda()

        actual = torch.empty_like(tensor, dtype=torch.half, device=torch.device('cuda'))

        decompress_int8(actual, tensor, scale)
        expected = (tensor.to(dtype=torch.float32) * scale[:, None]).to(dtype=torch.half)

        self.assertEqual(actual, expected)
//  [python]
def test_decompress_int8_cuda(self):
        tensor = torch.randint(-128, 127, (in_fea, out_fea), dtype=torch.int8).cuda()
        scale = torch.rand(in_fea).cuda()

        actual = torch.empty_like(tensor, dtype=torch.half, device=torch.device('cuda'))

        decompress_int8(actual, tensor, scale)
        expected = (tensor.to(dtype=torch.float32) * scale[:, None]).to(dtype=torch.half)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_decompress_int4_cuda(self):
        tensor = torch.randint(-8, 7, (in_fea, out_fea), dtype=torch.int8).cuda()
        scale = torch.rand(in_fea).cuda()

        packed = ((tensor[:, ::2] & 0b00001111) << 4) | (tensor[:, 1::2] & 0b00001111)

        actual = torch.empty_like(tensor, dtype=torch.half, device=torch.device('cuda'))

        decompress_int4(actual, packed, scale)
        expected = (
            torch.stack((packed >> 4, packed << 4 >> 4), -1).view(in_fea, -1).to(dtype=torch.float32) * scale[:, None]
        ).to(dtype=torch.half)

        self.assertEqual(actual, expected)
//  [python]
def test_decompress_int4_cuda(self):
        tensor = torch.randint(-8, 7, (in_fea, out_fea), dtype=torch.int8).cuda()
        scale = torch.rand(in_fea).cuda()

        packed = ((tensor[:, ::2] & 0b00001111) << 4) | (tensor[:, 1::2] & 0b00001111)

        actual = torch.empty_like(tensor, dtype=torch.half, device=torch.device('cuda'))

        decompress_int4(actual, packed, scale)
        expected = (
            torch.stack((packed >> 4, packed << 4 >> 4), -1).view(in_fea, -1).to(dtype=torch.float32) * scale[:, None]
        ).to(dtype=torch.half)

        torch.testing.assert_close(actual, expected)
//  [python]
def main(batch_size, num_tables, num_tasks, repeats) -> None:
    device = torch.device("cuda", 0)
    torch.cuda.set_device(device)
    hash_sizes = list(np.random.choice(range(50, 250), size=(num_tables)))
    lengths = []
    offsets = []
    indices = []
    for h in hash_sizes:
        l, o, i = generate_unary_feature(batch_size, h)
        lengths.append(torch.IntTensor(l).to(device))
        offsets.append(torch.IntTensor(o).to(device))
        indices.append(torch.IntTensor(i).to(device))
    lengths_tensor = torch.cat(lengths)
    indices_tensor = torch.cat(indices)
    offsets_tensor = torch.zeros(
        lengths_tensor.numel() + 1,
        dtype=lengths_tensor.dtype,
        device=lengths_tensor.device,
    )
    offsets_tensor[1:] = torch.ops.fbgemm.asynchronous_inclusive_cumsum(
        lengths_tensor.view(-1)
    )

    # forward
    ref_emb = MyModule(num_tasks, hash_sizes).to(device)
    unary_emb = batched_unary_embeddings_ops.BatchedUnaryEmbeddingBag(
        num_tasks, hash_sizes
    ).to(device)
    for i, param in enumerate(unary_emb.split_embedding_weights()):
        param.detach().copy_(ref_emb.emb_modules[i].weight)
    output_ref = ref_emb(offsets, indices)
    output = unary_emb(offsets_tensor, indices_tensor)
    t.assert_close(output_ref, output)
    # backward
    d_output = torch.randn([num_tasks, batch_size, len(hash_sizes)]).to(device) * 0.1
    output_ref.backward(d_output)
    output.backward(d_output)
    d_weight_ref = []
    for emb in ref_emb.emb_modules:
        d_weight_ref.append(emb.weight.grad)
    d_weight_ref = torch.cat(d_weight_ref).view(num_tasks, -1)
    d_weight = unary_emb.weight.grad
    # pyre-fixme[16]: Optional type has no attribute `squeeze`.
    t.assert_close(d_weight_ref, d_weight.squeeze())

    # A100 40MB L2 cache
    elapse, _ = benchmark_torch_function(ref_emb, (offsets, indices), iters=repeats)
    print("PyTorch EmbeddingBag forward", elapse)

    elapse, _ = benchmark_torch_function(
        unary_emb,
        (offsets_tensor, indices_tensor),
        iters=repeats,
    )
    print("Batched Unary Emb forward", elapse)

    output = ref_emb(offsets, indices)
    output.backward(d_output, retain_graph=True)
    elapse, _ = benchmark_torch_function(
        functools.partial(output.backward, retain_graph=True),
        (d_output,),
        iters=repeats,
    )
    print("PyTorch EmbeddingBag backward", elapse)

    output = unary_emb(offsets_tensor, indices_tensor)
    elapse, _ = benchmark_torch_function(
        functools.partial(output.backward, retain_graph=True),
        (d_output,),
        iters=repeats,
    )
    print("Batched Unary Emb backward", elapse)
//  [python]
def main(batch_size, num_tables, num_tasks, repeats) -> None:
    device = torch.device("cuda", 0)
    torch.cuda.set_device(device)
    hash_sizes = list(np.random.choice(range(50, 250), size=(num_tables)))
    lengths = []
    offsets = []
    indices = []
    for h in hash_sizes:
        l, o, i = generate_unary_feature(batch_size, h)
        lengths.append(torch.IntTensor(l).to(device))
        offsets.append(torch.IntTensor(o).to(device))
        indices.append(torch.IntTensor(i).to(device))
    lengths_tensor = torch.cat(lengths)
    indices_tensor = torch.cat(indices)
    offsets_tensor = torch.zeros(
        lengths_tensor.numel() + 1,
        dtype=lengths_tensor.dtype,
        device=lengths_tensor.device,
    )
    offsets_tensor[1:] = torch.ops.fbgemm.asynchronous_inclusive_cumsum(
        lengths_tensor.view(-1)
    )

    # forward
    ref_emb = MyModule(num_tasks, hash_sizes).to(device)
    unary_emb = batched_unary_embeddings_ops.BatchedUnaryEmbeddingBag(
        num_tasks, hash_sizes
    ).to(device)
    for i, param in enumerate(unary_emb.split_embedding_weights()):
        param.detach().copy_(ref_emb.emb_modules[i].weight)
    output_ref = ref_emb(offsets, indices)
    output = unary_emb(offsets_tensor, indices_tensor)
    torch.testing.assert_close(output_ref, output)
    # backward
    d_output = torch.randn([num_tasks, batch_size, len(hash_sizes)]).to(device) * 0.1
    output_ref.backward(d_output)
    output.backward(d_output)
    d_weight_ref = []
    for emb in ref_emb.emb_modules:
        d_weight_ref.append(emb.weight.grad)
    d_weight_ref = torch.cat(d_weight_ref).view(num_tasks, -1)
    d_weight = unary_emb.weight.grad
    # pyre-fixme[16]: Optional type has no attribute `squeeze`.
    torch.testing.assert_close(d_weight_ref, d_weight.squeeze())

    # A100 40MB L2 cache
    elapse, _ = benchmark_torch_function(ref_emb, (offsets, indices), iters=repeats)
    print("PyTorch EmbeddingBag forward", elapse)

    elapse, _ = benchmark_torch_function(
        unary_emb,
        (offsets_tensor, indices_tensor),
        iters=repeats,
    )
    print("Batched Unary Emb forward", elapse)

    output = ref_emb(offsets, indices)
    output.backward(d_output, retain_graph=True)
    elapse, _ = benchmark_torch_function(
        functools.partial(output.backward, retain_graph=True),
        (d_output,),
        iters=repeats,
    )
    print("PyTorch EmbeddingBag backward", elapse)

    output = unary_emb(offsets_tensor, indices_tensor)
    elapse, _ = benchmark_torch_function(
        functools.partial(output.backward, retain_graph=True),
        (d_output,),
        iters=repeats,
    )
    print("Batched Unary Emb backward", elapse)
//  [python]
def _test_baseline_and_shampoo(
        baseline_optim_factory: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        shampoo_optim_factory: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        device: torch.device,
    ) -> None:
        (
            baseline_params,
            baseline_loss,
        ) = DistributedShampooGraftingTest._train_quadratic(
            baseline_optim_factory,
            device=device,
        )
        shampoo_params, shampoo_loss = DistributedShampooGraftingTest._train_quadratic(
            shampoo_optim_factory,
            device=device,
        )
        torch.testing.assert_allclose(shampoo_loss, baseline_loss)
        torch.testing.assert_allclose(
            shampoo_params,
            baseline_params,
        )
//  [python]
def _test_baseline_and_shampoo(
        baseline_optim_factory: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        shampoo_optim_factory: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        device: torch.device,
    ) -> None:
        (
            baseline_params,
            baseline_loss,
        ) = DistributedShampooGraftingTest._train_quadratic(
            baseline_optim_factory,
            device=device,
        )
        shampoo_params, shampoo_loss = DistributedShampooGraftingTest._train_quadratic(
            shampoo_optim_factory,
            device=device,
        )
        torch.testing.assert_close(shampoo_loss, baseline_loss)
        torch.testing.assert_close(
            shampoo_params,
            baseline_params,
        )
//  [python]
print(f"Transcript: {beam_search_transcript_inc}")
print(f"WER: {beam_search_wer_inc}")

assert beam_search_result[0][0].words == beam_search_result_inc[0].words
assert beam_search_result[0][0].score == beam_search_result_inc[0].score
torch.testing.assert_allclose(beam_search_result[0][0].timesteps, beam_search_result_inc[0].timesteps)
//  [python]
print(f"Transcript: {beam_search_transcript_inc}")
print(f"WER: {beam_search_wer_inc}")

assert beam_search_result[0][0].words == beam_search_result_inc[0].words
assert beam_search_result[0][0].score == beam_search_result_inc[0].score
torch.testing.assert_close(beam_search_result[0][0].timesteps, beam_search_result_inc[0].timesteps)
//  [python]
def compare_rope(
    bsz: int, num_heads: int, embed_dim: int, seq_len: int, max_seq_len: int
) -> None:

    # make sure we have the right seed for generating outputs
    torch.manual_seed(0)

    head_dim = embed_dim // num_heads

    # generate input tensor
    x = torch.randn(bsz, seq_len, num_heads, head_dim)

    # Compute the reference tensors
    freq_cis = precompute_freqs_cis(dim=head_dim, end=max_seq_len * 2)
    x_out_ref = apply_rotary_emb(x, freqs_cis=freq_cis[:seq_len])

    # Compute the tensors from current implementation
    rope_emb = RotaryPositionalEmbeddings(dim=head_dim, max_seq_len=max_seq_len)
    x_out = rope_emb(x)

    # Validate correctness
    t.assert_close(x_out_ref, x_out, atol=1e-6, rtol=1e-5)

    # value: tensor(6.4543e-05)
    print(x_out.mean())

    # value: tensor(2165.7053)
    print(x_out.sum())

    # value: tensor(5.4546)
    print(x_out.max())

    curr_pos = 10
    x_out_ref = apply_rotary_emb(x, freqs_cis=freq_cis[curr_pos : curr_pos + seq_len])
    x_out = rope_emb(x, curr_pos=10)

    # Validate correctness
    t.assert_close(x_out_ref, x_out, atol=1e-6, rtol=1e-5)

    # value: tensor(0.0002)
    print(x_out.mean())

    # value: tensor(5158.3159)
    print(x_out.sum())

    # value: tensor(5.4543)
    print(x_out.max())
//  [python]
def compare_rope(
    bsz: int, num_heads: int, embed_dim: int, seq_len: int, max_seq_len: int
) -> None:

    # make sure we have the right seed for generating outputs
    torch.manual_seed(0)

    head_dim = embed_dim // num_heads

    # generate input tensor
    x = torch.randn(bsz, seq_len, num_heads, head_dim)

    # Compute the reference tensors
    freq_cis = precompute_freqs_cis(dim=head_dim, end=max_seq_len * 2)
    x_out_ref = apply_rotary_emb(x, freqs_cis=freq_cis[:seq_len])

    # Compute the tensors from current implementation
    rope_emb = RotaryPositionalEmbeddings(dim=head_dim, max_seq_len=max_seq_len)
    x_out = rope_emb(x)

    # Validate correctness
    torch.testing.assert_close(x_out_ref, x_out, atol=1e-6, rtol=1e-5)

    # value: tensor(6.4543e-05)
    print(x_out.mean())

    # value: tensor(2165.7053)
    print(x_out.sum())

    # value: tensor(5.4546)
    print(x_out.max())

    curr_pos = 10
    x_out_ref = apply_rotary_emb(x, freqs_cis=freq_cis[curr_pos : curr_pos + seq_len])
    x_out = rope_emb(x, curr_pos=10)

    # Validate correctness
    torch.testing.assert_close(x_out_ref, x_out, atol=1e-6, rtol=1e-5)

    # value: tensor(0.0002)
    print(x_out.mean())

    # value: tensor(5158.3159)
    print(x_out.sum())

    # value: tensor(5.4543)
    print(x_out.max())
//  [python]
def _test_accuracy_helper(
        self,
        labels: torch.Tensor,
        predictions: torch.Tensor,
        weights: torch.Tensor,
        expected_accuracy: torch.Tensor,
        threshold: float,
    ) -> None:
        num_task = labels.shape[0]
        batch_size = labels.shape[0]
        task_list = []
        inputs: Dict[str, Union[Dict[str, torch.Tensor], torch.Tensor]] = {
            "predictions": {},
            "labels": {},
            "weights": {},
        }
        for i in range(num_task):
            task_info = RecTaskInfo(
                name=f"Task:{i}",
                label_name="label",
                prediction_name="prediction",
                weight_name="weight",
            )
            task_list.append(task_info)
            # pyre-ignore
            inputs["predictions"][task_info.name] = predictions[i]
            # pyre-ignore
            inputs["labels"][task_info.name] = labels[i]
            # pyre-ignore
            inputs["weights"][task_info.name] = weights[i]

        accuracy = AccuracyMetric(
            world_size=1,
            my_rank=0,
            batch_size=batch_size,
            tasks=task_list,
            # pyre-ignore
            threshold=threshold,  # threshold is one of the kwargs
        )
        # pyre-ignore
        accuracy.update(**inputs)
        actual_accuracy = accuracy.compute()

        for task_id, task in enumerate(task_list):
            cur_actual_accuracy = actual_accuracy[
                f"accuracy-{task.name}|window_accuracy"
            ]
            cur_expected_accuracy = expected_accuracy[task_id].unsqueeze(dim=0)

            t.assert_allclose(
                cur_actual_accuracy,
                cur_expected_accuracy,
                atol=1e-4,
                rtol=1e-4,
                check_dtype=False,
                msg=f"Actual: {cur_actual_accuracy}, Expected: {cur_expected_accuracy}",
            )
//  [python]
def _test_accuracy_helper(
        self,
        labels: torch.Tensor,
        predictions: torch.Tensor,
        weights: torch.Tensor,
        expected_accuracy: torch.Tensor,
        threshold: float,
    ) -> None:
        num_task = labels.shape[0]
        batch_size = labels.shape[0]
        task_list = []
        inputs: Dict[str, Union[Dict[str, torch.Tensor], torch.Tensor]] = {
            "predictions": {},
            "labels": {},
            "weights": {},
        }
        for i in range(num_task):
            task_info = RecTaskInfo(
                name=f"Task:{i}",
                label_name="label",
                prediction_name="prediction",
                weight_name="weight",
            )
            task_list.append(task_info)
            # pyre-ignore
            inputs["predictions"][task_info.name] = predictions[i]
            # pyre-ignore
            inputs["labels"][task_info.name] = labels[i]
            # pyre-ignore
            inputs["weights"][task_info.name] = weights[i]

        accuracy = AccuracyMetric(
            world_size=1,
            my_rank=0,
            batch_size=batch_size,
            tasks=task_list,
            # pyre-ignore
            threshold=threshold,  # threshold is one of the kwargs
        )
        # pyre-ignore
        accuracy.update(**inputs)
        actual_accuracy = accuracy.compute()

        for task_id, task in enumerate(task_list):
            cur_actual_accuracy = actual_accuracy[
                f"accuracy-{task.name}|window_accuracy"
            ]
            cur_expected_accuracy = expected_accuracy[task_id].unsqueeze(dim=0)

            torch.testing.assert_close(
                cur_actual_accuracy,
                cur_expected_accuracy,
                atol=1e-4,
                rtol=1e-4,
                check_dtype=False,
                msg=f"Actual: {cur_actual_accuracy}, Expected: {cur_expected_accuracy}",
            )
//  [python]
def test_sht(grid):
    nlat = 16
    nlon = 2*nlat
    batch_size = 2
    if grid == "equiangular":
        mmax = nlat // 2
    else:
        mmax = nlat
    lmax = mmax
    norm = 'ortho'
    dtype = torch.float32

    sht_handle = SHT(dtype=dtype)

    # Create input
    coeffs = torch.zeros(batch_size, lmax, mmax, dtype=torch.complex64)
    coeffs[:, :lmax, :mmax] = torch.randn(batch_size, lmax, mmax, dtype=torch.complex64)
    
    signal = sht_handle.isht(coeffs, s=(nlat, nlon), grid=grid, norm=norm).to(torch.float32)

    coeffs = sht_handle.sht(signal, s=(lmax, mmax), grid=grid, norm=norm)
    rec = sht_handle.isht(coeffs, s=(nlat, nlon), grid=grid, norm=norm)
    t.assert_close(signal, rec, rtol=1e-4, atol=1e-4)
//  [python]
def test_sht(grid):
    nlat = 16
    nlon = 2*nlat
    batch_size = 2
    if grid == "equiangular":
        mmax = nlat // 2
    else:
        mmax = nlat
    lmax = mmax
    norm = 'ortho'
    dtype = torch.float32

    sht_handle = SHT(dtype=dtype)

    # Create input
    coeffs = torch.zeros(batch_size, lmax, mmax, dtype=torch.complex64)
    coeffs[:, :lmax, :mmax] = torch.randn(batch_size, lmax, mmax, dtype=torch.complex64)
    
    signal = sht_handle.isht(coeffs, s=(nlat, nlon), grid=grid, norm=norm).to(torch.float32)

    coeffs = sht_handle.sht(signal, s=(lmax, mmax), grid=grid, norm=norm)
    rec = sht_handle.isht(coeffs, s=(nlat, nlon), grid=grid, norm=norm)
    torch.testing.assert_close(signal, rec, rtol=1e-4, atol=1e-4)
//  [python]
def verify_mm_relu_mods(self, mm_only_mod, relu_mod):
        """
        Verify one module only does a mm op while the other
        performs both mm and relu ops in cascade
        """
        x = torch.randn(10, 5)
        t.assert_close(
            mm_only_mod(x), torch.mm(x, mm_only_mod.get_mul_matrix())
        )
        tracer = torch.fx.Tracer(param_shapes_constant=True)
        traced_graph = tracer.trace(mm_only_mod)

        # verify the graph module calculates the same result
        graph_mod_mm = torch.fx.GraphModule(mm_only_mod, traced_graph)
        t.assert_close(
            graph_mod_mm(x), torch.mm(x, mm_only_mod.get_mul_matrix())
        )

        # Make a new module with different parameter shape to go down the different
        # code path
        x = torch.randn(10, 15)
        t.assert_close(
            relu_mod(x), torch.relu(torch.mm(x, relu_mod.get_mul_matrix()))
        )

        tracer2 = torch.fx.Tracer(param_shapes_constant=True)
        traced_graph2 = tracer2.trace(relu_mod)

        # verify the graph module calculates the same result
        graph_mod_relu = torch.fx.GraphModule(relu_mod, traced_graph2)
        t.assert_close(
            graph_mod_relu(x), torch.relu(torch.mm(x, relu_mod.get_mul_matrix()))
        )

        graph1_node_targets = [n.target for n in traced_graph.nodes]
        graph2_node_targets = [n.target for n in traced_graph2.nodes]

        # the second graph has an exta relu function call node
        assert torch.mm in graph1_node_targets and torch.mm in graph2_node_targets
        assert (
            torch.relu not in graph1_node_targets and torch.relu in graph2_node_targets
        )
//  [python]
def verify_mm_relu_mods(self, mm_only_mod, relu_mod):
        """
        Verify one module only does a mm op while the other
        performs both mm and relu ops in cascade
        """
        x = torch.randn(10, 5)
        torch.testing.assert_close(
            mm_only_mod(x), torch.mm(x, mm_only_mod.get_mul_matrix())
        )
        tracer = torch.fx.Tracer(param_shapes_constant=True)
        traced_graph = tracer.trace(mm_only_mod)

        # verify the graph module calculates the same result
        graph_mod_mm = torch.fx.GraphModule(mm_only_mod, traced_graph)
        torch.testing.assert_close(
            graph_mod_mm(x), torch.mm(x, mm_only_mod.get_mul_matrix())
        )

        # Make a new module with different parameter shape to go down the different
        # code path
        x = torch.randn(10, 15)
        torch.testing.assert_close(
            relu_mod(x), torch.relu(torch.mm(x, relu_mod.get_mul_matrix()))
        )

        tracer2 = torch.fx.Tracer(param_shapes_constant=True)
        traced_graph2 = tracer2.trace(relu_mod)

        # verify the graph module calculates the same result
        graph_mod_relu = torch.fx.GraphModule(relu_mod, traced_graph2)
        torch.testing.assert_close(
            graph_mod_relu(x), torch.relu(torch.mm(x, relu_mod.get_mul_matrix()))
        )

        graph1_node_targets = [n.target for n in traced_graph.nodes]
        graph2_node_targets = [n.target for n in traced_graph2.nodes]

        # the second graph has an exta relu function call node
        assert torch.mm in graph1_node_targets and torch.mm in graph2_node_targets
        assert (
            torch.relu not in graph1_node_targets and torch.relu in graph2_node_targets
        )
//  [python]
def test_meta_access(self):
        dataset = SqlIndexDataset(
            sqlite_metadata_file=METADATA_FILE,
            remove_empty_masks=False,
            subset_lists_file=SET_LIST_FILE,
            subsets=["train"],
            frame_data_builder_FrameDataBuilder_args=NO_BLOBS_KWARGS,
        )

        self.assertEqual(len(dataset), 50)

        for idx in [10, ("cat0_seq2", 2)]:
            example_meta = dataset.meta[idx]
            example = dataset[idx]
            self.assertEqual(example_meta.sequence_name, example.sequence_name)
            self.assertEqual(example_meta.frame_number, example.frame_number)
            self.assertEqual(example_meta.frame_timestamp, example.frame_timestamp)
            self.assertEqual(example_meta.sequence_category, example.sequence_category)
            self.assertTrue(torch.allclose(example_meta.camera.R, example.camera.R))
            self.assertTrue(torch.allclose(example_meta.camera.T, example.camera.T))
            self.assertTrue(
                torch.allclose(example_meta.camera.focal_length, example.camera.focal_length)
            )
            self.assertTrue(
                torch.allclose(example_meta.camera.principal_point, example.camera.principal_point)
            )
//  [python]
def test_meta_access(self):
        dataset = SqlIndexDataset(
            sqlite_metadata_file=METADATA_FILE,
            remove_empty_masks=False,
            subset_lists_file=SET_LIST_FILE,
            subsets=["train"],
            frame_data_builder_FrameDataBuilder_args=NO_BLOBS_KWARGS,
        )

        self.assertEqual(len(dataset), 50)

        for idx in [10, ("cat0_seq2", 2)]:
            example_meta = dataset.meta[idx]
            example = dataset[idx]
            self.assertEqual(example_meta.sequence_name, example.sequence_name)
            self.assertEqual(example_meta.frame_number, example.frame_number)
            self.assertEqual(example_meta.frame_timestamp, example.frame_timestamp)
            self.assertEqual(example_meta.sequence_category, example.sequence_category)
            torch.testing.assert_close(example_meta.camera.R, example.camera.R)
            torch.testing.assert_close(example_meta.camera.T, example.camera.T)
            torch.testing.assert_close(
                example_meta.camera.focal_length, example.camera.focal_length
            )
            torch.testing.assert_close(
                example_meta.camera.principal_point, example.camera.principal_point
            )
//  [python]
def test_loss(self, tmpdir, monkeypatch):
        ckpt = "small_test_ckpt_tune"
        ckpt_path = Path(CKPT_MODEL_PATHS[ckpt])
        ckpt_dir = ckpt_path.parent
        log_file = gen_log_file_name(tmpdir)
        cmd = f"""
        tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed
            --config llama2/7B_lora \
            output_dir={tmpdir} \
            checkpointer=torchtune.utils.FullModelTorchTuneCheckpointer \
            checkpointer.checkpoint_dir='{ckpt_dir}' \
            checkpointer.checkpoint_files=[{ckpt_path}]\
            checkpointer.output_dir={tmpdir} \
            checkpointer.model_type=LLAMA2 \
            metric_logger.filename={log_file} \
        """.split()

        model_config = lora_llama2_test_config(
            lora_attn_modules=["q_proj", "k_proj", "v_proj", "output_proj"],
            apply_lora_to_mlp=False,
            apply_lora_to_output=False,
            lora_rank=8,
            lora_alpha=16,
        )

        cmd = cmd + self._get_test_config_overrides() + model_config
        monkeypatch.setattr(sys, "argv", cmd)
        runpy.run_path(TUNE_PATH, run_name="__main__")
        loss_values = get_loss_values_from_metric_logger(log_file)
        expected_loss_values = self._fetch_expected_loss_values()
        t.assert_allclose(
            loss_values, expected_loss_values, rtol=1e-5, atol=1e-5
        )
//  [python]
def test_loss(self, tmpdir, monkeypatch):
        ckpt = "small_test_ckpt_tune"
        ckpt_path = Path(CKPT_MODEL_PATHS[ckpt])
        ckpt_dir = ckpt_path.parent
        log_file = gen_log_file_name(tmpdir)
        cmd = f"""
        tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed
            --config llama2/7B_lora \
            output_dir={tmpdir} \
            checkpointer=torchtune.utils.FullModelTorchTuneCheckpointer \
            checkpointer.checkpoint_dir='{ckpt_dir}' \
            checkpointer.checkpoint_files=[{ckpt_path}]\
            checkpointer.output_dir={tmpdir} \
            checkpointer.model_type=LLAMA2 \
            metric_logger.filename={log_file} \
        """.split()

        model_config = lora_llama2_test_config(
            lora_attn_modules=["q_proj", "k_proj", "v_proj", "output_proj"],
            apply_lora_to_mlp=False,
            apply_lora_to_output=False,
            lora_rank=8,
            lora_alpha=16,
        )

        cmd = cmd + self._get_test_config_overrides() + model_config
        monkeypatch.setattr(sys, "argv", cmd)
        runpy.run_path(TUNE_PATH, run_name="__main__")
        loss_values = get_loss_values_from_metric_logger(log_file)
        expected_loss_values = self._fetch_expected_loss_values()
        torch.testing.assert_close(
            loss_values, expected_loss_values, rtol=1e-5, atol=1e-5
        )
//  [python]
def test_SphericalConv(factorization, implementation):
    """Test for SphericalConv (2D only)
    
    Compares Factorized and Dense convolution output
    Verifies that a dense conv and factorized conv with the same weight produce the same output

    Checks the output size

    Verifies that dynamically changing the number of Fourier modes doesn't break the conv
    """
    n_modes = (6, 6)

    conv = SphericalConv(
        3, 3, n_modes, n_layers=1, bias=False, implementation=implementation, factorization=factorization)

    conv_dense = SphericalConv(
        3, 3, n_modes, n_layers=1, bias=False, implementation='reconstructed', factorization=None)

    conv_dense.weight[0] = FactorizedTensor.from_tensor(conv.weight[0].to_tensor(), rank=None, factorization='ComplexDense')
    x = torch.randn(2, 3, *(12, 12))

    res_dense = conv_dense(x)
    res = conv(x)

    torch.testing.assert_allclose(res_dense, res)

    # Downsample outputs
    block = SphericalConv(
        3, 4, n_modes, n_layers=1, output_scaling_factor=0.5)

    x = torch.randn(2, 3, *(12, 12))
    res = block(x)
    assert(list(res.shape[2:]) == [12//2, 12//2])

    # Upsample outputs
    block = SphericalConv(
        3, 4, n_modes, n_layers=1, output_scaling_factor=2)

    x = torch.randn(2, 3, *(12, 12))
    res = block(x)
    assert res.shape[1] == 4 # Check out channels
    assert(list(res.shape[2:]) == [12*2, 12*2])

    # Test change of grid
    block = SphericalConv(
        4, 4, n_modes, n_layers=2, sht_grids=["equiangular", "legendre-gauss", "equiangular"])
    x = torch.randn(2, 4, *(12, 12))
    res = block[0](x)
    res = block[1](res)
    assert(res.shape[2:] == x.shape[2:])

    res = block[0].transform(x)
    res = block[1].transform(res)
    assert(res.shape[2:] == x.shape[2:])
//  [python]
def test_SphericalConv(factorization, implementation):
    """Test for SphericalConv (2D only)
    
    Compares Factorized and Dense convolution output
    Verifies that a dense conv and factorized conv with the same weight produce the same output

    Checks the output size

    Verifies that dynamically changing the number of Fourier modes doesn't break the conv
    """
    n_modes = (6, 6)

    conv = SphericalConv(
        3, 3, n_modes, n_layers=1, bias=False, implementation=implementation, factorization=factorization)

    conv_dense = SphericalConv(
        3, 3, n_modes, n_layers=1, bias=False, implementation='reconstructed', factorization=None)

    conv_dense.weight[0] = FactorizedTensor.from_tensor(conv.weight[0].to_tensor(), rank=None, factorization='ComplexDense')
    x = torch.randn(2, 3, *(12, 12))

    res_dense = conv_dense(x)
    res = conv(x)

    torch.testing.assert_close(res_dense, res)

    # Downsample outputs
    block = SphericalConv(
        3, 4, n_modes, n_layers=1, output_scaling_factor=0.5)

    x = torch.randn(2, 3, *(12, 12))
    res = block(x)
    assert(list(res.shape[2:]) == [12//2, 12//2])

    # Upsample outputs
    block = SphericalConv(
        3, 4, n_modes, n_layers=1, output_scaling_factor=2)

    x = torch.randn(2, 3, *(12, 12))
    res = block(x)
    assert res.shape[1] == 4 # Check out channels
    assert(list(res.shape[2:]) == [12*2, 12*2])

    # Test change of grid
    block = SphericalConv(
        4, 4, n_modes, n_layers=2, sht_grids=["equiangular", "legendre-gauss", "equiangular"])
    x = torch.randn(2, 4, *(12, 12))
    res = block[0](x)
    res = block[1](res)
    assert(res.shape[2:] == x.shape[2:])

    res = block[0].transform(x)
    res = block[1].transform(res)
    assert(res.shape[2:] == x.shape[2:])