jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_standard_transposed_conv_cpu():
    '''Simply compares our result to PyTorch's implementation'''
    input = torch.ones(params.bs, params.in_channels, 5, 5).double()
    input.requires_grad = True

    pytorch_layer = torch.nn.ConvTranspose2d(params.in_channels,
                                             params.out_channels,
                                             kernel_size=3,
                                             stride=2,
                                             padding=1,
                                             dilation=2).double()
    my_layer = TransposedConvolution(params.in_channels,
                                     params.out_channels,
                                     kernel_size=3,
                                     stride=2,
                                     padding=1,
                                     dilation=2).double()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    torch.testing.assert_close(my_output, pytorch_output)
//  [python]
def test_standard_transposed_conv_cuda():
    '''Simply compares our result to PyTorch's implementation'''

    input = torch.ones(params.bs, params.in_channels, 5, 5).double().cuda()
    input.requires_grad = True

    pytorch_layer = torch.nn.ConvTranspose2d(params.in_channels,
                                             params.out_channels,
                                             kernel_size=3,
                                             stride=2,
                                             padding=1,
                                             dilation=2).double().cuda()
    my_layer = TransposedConvolution(params.in_channels,
                                     params.out_channels,
                                     kernel_size=3,
                                     stride=2,
                                     padding=1,
                                     dilation=2).double().cuda()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    testing.assert_allclose(my_output, pytorch_output)
//  [python]
def test_standard_transposed_conv_cuda():
    '''Simply compares our result to PyTorch's implementation'''

    input = torch.ones(params.bs, params.in_channels, 5, 5).double().cuda()
    input.requires_grad = True

    pytorch_layer = torch.nn.ConvTranspose2d(params.in_channels,
                                             params.out_channels,
                                             kernel_size=3,
                                             stride=2,
                                             padding=1,
                                             dilation=2).double().cuda()
    my_layer = TransposedConvolution(params.in_channels,
                                     params.out_channels,
                                     kernel_size=3,
                                     stride=2,
                                     padding=1,
                                     dilation=2).double().cuda()

    # Set the layer weights
    pytorch_layer.weight.data.fill_(1)
    pytorch_layer.bias.data.fill_(2)
    my_layer.weight.data.fill_(1)
    my_layer.bias.data.fill_(2)

    # Run a forward pass
    pytorch_output = pytorch_layer(input)
    my_output = my_layer(input)

    # Run a numerical gradient check
    gradcheck_res = gradcheck(my_layer, (input))

    # Ensure our implementation passes gradcheck
    assert gradcheck_res

    # Make sure they match
    torch.testing.assert_close(my_output, pytorch_output)
//  [python]
def test_correctness(self):
        x = torch.randn(32, 4)
        dist = Delta(v=x, event_dim=1)

        s = dist.sample((10,))
        torch.testing.assert_allclose(s, x.expand(10, 32, 4))
//  [python]
def test_correctness(self):
        x = torch.randn(32, 4)
        dist = Delta(v=x, event_dim=1)

        s = dist.sample((10,))
        torch.testing.assert_close(s, x.expand(10, 32, 4))
//  [python]
def test_entropy(self):
        x = torch.randn(32, 4)
        dist = Delta(v=x)
        torch.testing.assert_allclose(dist.entropy(), torch.zeros(32, 4))
        torch.testing.assert_allclose(dist.variance, torch.zeros(32, 4))

        dist = Delta(v=x, event_dim=1)
        torch.testing.assert_allclose(dist.entropy(), torch.zeros(32))
        torch.testing.assert_allclose(dist.variance, torch.zeros(32, 4))
//  [python]
def test_entropy(self):
        x = torch.randn(32, 4)
        dist = Delta(v=x)
        torch.testing.assert_close(dist.entropy(), torch.zeros(32, 4))
        torch.testing.assert_close(dist.variance, torch.zeros(32, 4))

        dist = Delta(v=x, event_dim=1)
        torch.testing.assert_close(dist.entropy(), torch.zeros(32))
        torch.testing.assert_close(dist.variance, torch.zeros(32, 4))
//  [python]
def test_strided_waveform_batch_streaming_snip_edges_false(deterministic_rng):
    x = torch.arange(16000).unsqueeze(0)
    window_length = 400
    window_shift = 160

    # reference offline forward pass
    y = _get_strided_batch(
        waveform=x,
        window_length=window_length,
        window_shift=window_shift,
        snip_edges=False,
    )
    assert y.shape == torch.Size([1, 100, window_length])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        y_chunk, remainder = _get_strided_batch_streaming(
            x_chunk, window_length=400, window_shift=160, prev_remainder=remainder
        )
        frames.append(y_chunk)

    # one last extra iteration to pump out the frames that the offline version gets due to padding
    # note that our implementation will flip the last N samples so we do the same for identical results
    x_chunk = torch.flip(x[:, -window_shift:], dims=(1,))
    y_chunk, remainder = _get_strided_batch_streaming(
        x_chunk,
        window_length=window_length,
        window_shift=window_shift,
        prev_remainder=remainder,
        snip_edges=False,
    )
    frames.append(y_chunk)

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_allclose(y_online, y)
//  [python]
def test_strided_waveform_batch_streaming_snip_edges_false(deterministic_rng):
    x = torch.arange(16000).unsqueeze(0)
    window_length = 400
    window_shift = 160

    # reference offline forward pass
    y = _get_strided_batch(
        waveform=x,
        window_length=window_length,
        window_shift=window_shift,
        snip_edges=False,
    )
    assert y.shape == torch.Size([1, 100, window_length])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        y_chunk, remainder = _get_strided_batch_streaming(
            x_chunk, window_length=400, window_shift=160, prev_remainder=remainder
        )
        frames.append(y_chunk)

    # one last extra iteration to pump out the frames that the offline version gets due to padding
    # note that our implementation will flip the last N samples so we do the same for identical results
    x_chunk = torch.flip(x[:, -window_shift:], dims=(1,))
    y_chunk, remainder = _get_strided_batch_streaming(
        x_chunk,
        window_length=window_length,
        window_shift=window_shift,
        prev_remainder=remainder,
        snip_edges=False,
    )
    frames.append(y_chunk)

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_close(y_online, y)
//  [python]
def test_strided_waveform_batch_streaming_snip_edges_true(deterministic_rng):
    x = torch.arange(16000).unsqueeze(0)
    window_length = 400
    window_shift = 160

    # reference offline forward pass
    y = _get_strided_batch(
        waveform=x,
        window_length=window_length,
        window_shift=window_shift,
        snip_edges=True,
    )
    assert y.shape == torch.Size([1, 98, window_length])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        y_chunk, remainder = _get_strided_batch_streaming(
            x_chunk,
            window_length=400,
            window_shift=160,
            prev_remainder=remainder,
            snip_edges=True,
        )
        frames.append(y_chunk)

    # Note: unlike the case of snip_edges=False,
    # we don't run one last step with reflected samples here
    # to match the behavior of offline snip_edges=True.

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_allclose(y_online, y)
//  [python]
def test_strided_waveform_batch_streaming_snip_edges_true(deterministic_rng):
    x = torch.arange(16000).unsqueeze(0)
    window_length = 400
    window_shift = 160

    # reference offline forward pass
    y = _get_strided_batch(
        waveform=x,
        window_length=window_length,
        window_shift=window_shift,
        snip_edges=True,
    )
    assert y.shape == torch.Size([1, 98, window_length])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        y_chunk, remainder = _get_strided_batch_streaming(
            x_chunk,
            window_length=400,
            window_shift=160,
            prev_remainder=remainder,
            snip_edges=True,
        )
        frames.append(y_chunk)

    # Note: unlike the case of snip_edges=False,
    # we don't run one last step with reflected samples here
    # to match the behavior of offline snip_edges=True.

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_close(y_online, y)
//  [python]
def test_wav2win_streaming(deterministic_rng):
    x = torch.randn(1, 16000, dtype=torch.float32)
    t = Wav2Win()
    window_length = 400
    window_shift = 160
    assert t._length == window_length
    assert t._shift == window_shift

    # reference offline forward pass
    y, _ = t(x)
    assert y.shape == torch.Size([1, 100, window_length])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        (y_chunk, _), remainder = t.online_inference(x_chunk, context=remainder)
        frames.append(y_chunk)

    # one last extra iteration to pump out the frames that the offline version gets due to padding
    x_chunk = torch.flip(x[:, -window_shift:], dims=(1,))
    (y_chunk, _), remainder = t.online_inference(x_chunk, context=remainder)
    frames.append(y_chunk)

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_allclose(y_online, y)
//  [python]
def test_wav2win_streaming(deterministic_rng):
    x = torch.randn(1, 16000, dtype=torch.float32)
    t = Wav2Win()
    window_length = 400
    window_shift = 160
    assert t._length == window_length
    assert t._shift == window_shift

    # reference offline forward pass
    y, _ = t(x)
    assert y.shape == torch.Size([1, 100, window_length])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        (y_chunk, _), remainder = t.online_inference(x_chunk, context=remainder)
        frames.append(y_chunk)

    # one last extra iteration to pump out the frames that the offline version gets due to padding
    x_chunk = torch.flip(x[:, -window_shift:], dims=(1,))
    (y_chunk, _), remainder = t.online_inference(x_chunk, context=remainder)
    frames.append(y_chunk)

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_close(y_online, y)
//  [python]
def test_wav2logfilterbank_streaming(deterministic_rng, layer_type, feat_dim):
    x = torch.randn(1, 16000, dtype=torch.float32)
    t = layer_type()
    window_length = 400
    window_shift = 160
    assert t.wav2win._length == window_length
    assert t.wav2win._shift == window_shift

    # reference offline forward pass
    y = t(x)
    assert y.shape == torch.Size([1, 100, feat_dim])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        y_chunk, remainder = t.online_inference(x_chunk, context=remainder)
        frames.append(y_chunk)

    # one last extra iteration to pump out the frames that the offline version gets due to padding
    x_chunk = torch.flip(x[:, -window_shift:], dims=(1,))
    y_chunk, remainder = t.online_inference(x_chunk, context=remainder)
    frames.append(y_chunk)

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_allclose(y_online, y)
//  [python]
def test_wav2logfilterbank_streaming(deterministic_rng, layer_type, feat_dim):
    x = torch.randn(1, 16000, dtype=torch.float32)
    t = layer_type()
    window_length = 400
    window_shift = 160
    assert t.wav2win._length == window_length
    assert t.wav2win._shift == window_shift

    # reference offline forward pass
    y = t(x)
    assert y.shape == torch.Size([1, 100, feat_dim])

    # online pass under test
    frames = []
    chunk_size = 1200
    num_chunks = math.ceil(x.size(1) / chunk_size)
    remainder = None
    for i in range(num_chunks):
        x_chunk = x[:, i * chunk_size : (i + 1) * chunk_size]
        y_chunk, remainder = t.online_inference(x_chunk, context=remainder)
        frames.append(y_chunk)

    # one last extra iteration to pump out the frames that the offline version gets due to padding
    x_chunk = torch.flip(x[:, -window_shift:], dims=(1,))
    y_chunk, remainder = t.online_inference(x_chunk, context=remainder)
    frames.append(y_chunk)

    y_online = torch.cat(frames, dim=1)

    assert y.shape == y_online.shape

    torch.testing.assert_close(y_online, y)
//  [python]
def test_multihead_selfattn(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        fake_sample = torch.rand(2, 10, 8)
        fake_sample[1, 3:, :] = 0
        fake_lengths = torch.LongTensor([10, 3])
        padding_mask = lengths_to_padding_mask(fake_lengths)
        fake_sample2 = fake_sample[1:, :3, :]
        padding_mask2 = lengths_to_padding_mask(fake_lengths[1].unsqueeze(0))
        attn = HyenaOperator(8, 10, num_heads=4)
        attn.eval()
        attn_out = attn(fake_sample, padding_mask)
        attn_out2 = attn(fake_sample2, padding_mask2)
        torch.testing.assert_allclose(attn_out[1, :3, :], attn_out2[0])
        self.assertTrue(
            torch.all(attn_out[1, 3:, :] == 0.0), f"non-zero entries in {attn_out[1, 3:, :]}")
//  [python]
def test_multihead_selfattn(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        fake_sample = torch.rand(2, 10, 8)
        fake_sample[1, 3:, :] = 0
        fake_lengths = torch.LongTensor([10, 3])
        padding_mask = lengths_to_padding_mask(fake_lengths)
        fake_sample2 = fake_sample[1:, :3, :]
        padding_mask2 = lengths_to_padding_mask(fake_lengths[1].unsqueeze(0))
        attn = HyenaOperator(8, 10, num_heads=4)
        attn.eval()
        attn_out = attn(fake_sample, padding_mask)
        attn_out2 = attn(fake_sample2, padding_mask2)
        torch.testing.assert_close(attn_out[1, :3, :], attn_out2[0])
        self.assertTrue(
            torch.all(attn_out[1, 3:, :] == 0.0), f"non-zero entries in {attn_out[1, 3:, :]}")
//  [python]
def test_encoder_batch(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        batchnorm_args.input_feat_per_channel = 8
        batchnorm_args.encoder_layers = 1
        fake_sample = torch.rand(5, 27, 8)
        fake_sample[1, 13:, :] = 0
        fake_sample[2, 8:, :] = 0
        fake_sample[3, 8:, :] = 0
        fake_sample[4, 5:, :] = 0
        fake_lengths = torch.LongTensor([27, 13, 8, 8, 5])
        encoder = ConformerHyenaEncoder(batchnorm_args, self.fake_dict)
        encoder.eval()
        net_out = encoder.forward(fake_sample, fake_lengths, return_all_hiddens=True)

        def test_item(item_idx):
            item_len = fake_lengths[item_idx].item()
            fake_sample2 = fake_sample[item_idx, :item_len, :]
            net_out2 = encoder.forward(
                fake_sample2.unsqueeze(0), fake_lengths[item_idx].unsqueeze(0), return_all_hiddens=True)
            torch.testing.assert_allclose(
                    net_out["encoder_out"][0][:item_len, item_idx, :],
                    net_out2["encoder_out"][0][:, 0, :])

        for i in range(5):
            test_item(i)
//  [python]
def test_encoder_batch(self):
        batchnorm_args = copy.deepcopy(self.base_args)
        batchnorm_args.no_syncbatchnorm = True
        batchnorm_args.encoder_embed_dim = 8
        batchnorm_args.input_feat_per_channel = 8
        batchnorm_args.encoder_layers = 1
        fake_sample = torch.rand(5, 27, 8)
        fake_sample[1, 13:, :] = 0
        fake_sample[2, 8:, :] = 0
        fake_sample[3, 8:, :] = 0
        fake_sample[4, 5:, :] = 0
        fake_lengths = torch.LongTensor([27, 13, 8, 8, 5])
        encoder = ConformerHyenaEncoder(batchnorm_args, self.fake_dict)
        encoder.eval()
        net_out = encoder.forward(fake_sample, fake_lengths, return_all_hiddens=True)

        def test_item(item_idx):
            item_len = fake_lengths[item_idx].item()
            fake_sample2 = fake_sample[item_idx, :item_len, :]
            net_out2 = encoder.forward(
                fake_sample2.unsqueeze(0), fake_lengths[item_idx].unsqueeze(0), return_all_hiddens=True)
            torch.testing.assert_close(
                    net_out["encoder_out"][0][:item_len, item_idx, :],
                    net_out2["encoder_out"][0][:, 0, :])

        for i in range(5):
            test_item(i)
//  [python]
def test_interpolate_points_scriptability(self) -> None:
        """
        Tests that the function interpolate_points scripts properly.
        """

        class tmp_module(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()

            def forward(self, coords: torch.Tensor, max_points: int, interpolation: str) -> torch.Tensor:
                result = interpolate_points(coords, max_points, interpolation)
                return result

        to_script = tmp_module()
        scripted = torch.jit.script(to_script)

        test_coords = torch.tensor([[1, 1], [3, 1], [5, 1]], dtype=torch.float64)

        py_result = to_script.forward(test_coords, 5, 'linear')
        script_result = scripted.forward(test_coords, 5, 'linear')

        torch.testing.assert_allclose(py_result, script_result)
//  [python]
def test_interpolate_points_scriptability(self) -> None:
        """
        Tests that the function interpolate_points scripts properly.
        """

        class tmp_module(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()

            def forward(self, coords: torch.Tensor, max_points: int, interpolation: str) -> torch.Tensor:
                result = interpolate_points(coords, max_points, interpolation)
                return result

        to_script = tmp_module()
        scripted = torch.jit.script(to_script)

        test_coords = torch.tensor([[1, 1], [3, 1], [5, 1]], dtype=torch.float64)

        py_result = to_script.forward(test_coords, 5, 'linear')
        script_result = scripted.forward(test_coords, 5, 'linear')

        torch.testing.assert_close(py_result, script_result)
//  [python]
def test_voxelize(
    grid: VoxelGrid,
    coordinates_m: Tensor,
    indices_: Tensor,
    values_: Tensor,
    counts_: Tensor,
) -> None:
    """Test voxelizing a point cloud."""
    indices, values, counts = voxelize(
        coordinates_m,
        coordinates_m,
        list(grid.min_world_coordinates_m),
        list(grid.delta_m_per_cell),
        list(grid.grid_size),
        cluster_type=ClusterType.MEAN,
    )

    torch.testing.assert_allclose(indices, indices_)
    torch.testing.assert_allclose(values, values_)
    torch.testing.assert_allclose(counts, counts_)
//  [python]
def test_voxelize(
    grid: VoxelGrid,
    coordinates_m: Tensor,
    indices_: Tensor,
    values_: Tensor,
    counts_: Tensor,
) -> None:
    """Test voxelizing a point cloud."""
    indices, values, counts = voxelize(
        coordinates_m,
        coordinates_m,
        list(grid.min_world_coordinates_m),
        list(grid.delta_m_per_cell),
        list(grid.grid_size),
        cluster_type=ClusterType.MEAN,
    )

    torch.testing.assert_close(indices, indices_)
    torch.testing.assert_close(values, values_)
    torch.testing.assert_close(counts, counts_)
//  [python]
def test_expectation():
    # Test E[X]
    μ = torch.tensor(1.)
    σ2 = torch.tensor(2.)

    expected = μ

    expectation = normal(lambda x: x)
    actual = expectation(μ, σ2)

    torch.testing.assert_allclose(expected, actual)
//  [python]
def test_expectation():
    # Test E[X]
    μ = torch.tensor(1.)
    σ2 = torch.tensor(2.)

    expected = μ

    expectation = normal(lambda x: x)
    actual = expectation(μ, σ2)

    torch.testing.assert_close(expected, actual)
//  [python]
def test_variance():
    # Test E[X**2]
    μ = torch.tensor(1.)
    σ2 = torch.tensor(2.)

    expected = μ**2 + σ2

    expectation = normal(lambda x: x**2)
    actual = expectation(μ, σ2)

    torch.testing.assert_allclose(expected, actual)
//  [python]
def test_variance():
    # Test E[X**2]
    μ = torch.tensor(1.)
    σ2 = torch.tensor(2.)

    expected = μ**2 + σ2

    expectation = normal(lambda x: x**2)
    actual = expectation(μ, σ2)

    torch.testing.assert_close(expected, actual)
//  [python]
def test_voxelize_concatenate_kernel(
    grid: RegularGrid,
    coordinates_m: Tensor,
    indices_: Tensor,
    values_: Tensor,
    counts_: Tensor,
) -> None:
    """Test voxelization with truncated reduction."""
    indices, values, counts = voxelize(
        coordinates_m,
        coordinates_m,
        list(grid.min_world_coordinates_m),
        list(grid.delta_m_per_cell),
        list(grid.grid_size),
        cluster_type=ClusterType.CONCATENATE,
        max_num_values=2,
    )

    torch.testing.assert_allclose(indices, indices_)
    torch.testing.assert_allclose(values, values_)
    torch.testing.assert_allclose(counts, counts_)
//  [python]
def test_voxelize_concatenate_kernel(
    grid: RegularGrid,
    coordinates_m: Tensor,
    indices_: Tensor,
    values_: Tensor,
    counts_: Tensor,
) -> None:
    """Test voxelization with truncated reduction."""
    indices, values, counts = voxelize(
        coordinates_m,
        coordinates_m,
        list(grid.min_world_coordinates_m),
        list(grid.delta_m_per_cell),
        list(grid.grid_size),
        cluster_type=ClusterType.CONCATENATE,
        max_num_values=2,
    )

    torch.testing.assert_close(indices, indices_)
    torch.testing.assert_close(values, values_)
    torch.testing.assert_close(counts, counts_)
//  [python]
def test_interpolate_points_functionality(self) -> None:
        """
        Test interpolating coordinate points.
        """
        coords = torch.tensor([[1, 1], [3, 1], [5, 1]], dtype=torch.float64)

        interpolated_coords = interpolate_points(coords, 5, interpolation='linear')
        self.assertEqual(interpolated_coords.shape, (5, 2))
        torch.testing.assert_allclose(coords, interpolated_coords[::2])
        torch.testing.assert_allclose(interpolated_coords[:, 1], torch.ones((5), dtype=torch.float64))
        self.assertTrue(interpolated_coords[1][0].item() > interpolated_coords[0][0].item())
        self.assertTrue(interpolated_coords[1][0].item() < interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() > interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() < interpolated_coords[4][0].item())

        interpolated_coords = interpolate_points(coords, 5, interpolation='area')
        self.assertEqual(interpolated_coords.shape, (5, 2))
        torch.testing.assert_allclose(coords, interpolated_coords[::2])
        torch.testing.assert_allclose(interpolated_coords[:, 1], torch.ones((5), dtype=torch.float64))
        self.assertTrue(interpolated_coords[1][0].item() > interpolated_coords[0][0].item())
        self.assertTrue(interpolated_coords[1][0].item() < interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() > interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() < interpolated_coords[4][0].item())
//  [python]
def test_interpolate_points_functionality(self) -> None:
        """
        Test interpolating coordinate points.
        """
        coords = torch.tensor([[1, 1], [3, 1], [5, 1]], dtype=torch.float64)

        interpolated_coords = interpolate_points(coords, 5, interpolation='linear')
        self.assertEqual(interpolated_coords.shape, (5, 2))
        torch.testing.assert_close(coords, interpolated_coords[::2])
        torch.testing.assert_close(interpolated_coords[:, 1], torch.ones((5), dtype=torch.float64))
        self.assertTrue(interpolated_coords[1][0].item() > interpolated_coords[0][0].item())
        self.assertTrue(interpolated_coords[1][0].item() < interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() > interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() < interpolated_coords[4][0].item())

        interpolated_coords = interpolate_points(coords, 5, interpolation='area')
        self.assertEqual(interpolated_coords.shape, (5, 2))
        torch.testing.assert_close(coords, interpolated_coords[::2])
        torch.testing.assert_close(interpolated_coords[:, 1], torch.ones((5), dtype=torch.float64))
        self.assertTrue(interpolated_coords[1][0].item() > interpolated_coords[0][0].item())
        self.assertTrue(interpolated_coords[1][0].item() < interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() > interpolated_coords[2][0].item())
        self.assertTrue(interpolated_coords[3][0].item() < interpolated_coords[4][0].item())
//  [python]
def test_c10_layer_norm(self):
        # test that we can call c10 ops and they return a reasonable result
        X = torch.rand(5, 5, dtype=torch.float)
        weight = torch.rand(*X.size()[1:], dtype=torch.float)
        bias = torch.rand(*X.size()[1:], dtype=torch.float)
        epsilon = 1e-4

        expected_norm = torch.nn.functional.layer_norm(
            X, X.size()[1:], weight=weight, bias=bias, eps=epsilon)
        actual_norm, actual_mean, actual_stdev = \
            torch.ops._caffe2.LayerNorm(torch.tensor(X), torch.tensor(
                weight), torch.tensor(bias), 1, epsilon, True)
        torch.testing.assert_allclose(expected_norm, actual_norm)
//  [python]
def test_c10_layer_norm(self):
        # test that we can call c10 ops and they return a reasonable result
        X = torch.rand(5, 5, dtype=torch.float)
        weight = torch.rand(*X.size()[1:], dtype=torch.float)
        bias = torch.rand(*X.size()[1:], dtype=torch.float)
        epsilon = 1e-4

        expected_norm = torch.nn.functional.layer_norm(
            X, X.size()[1:], weight=weight, bias=bias, eps=epsilon)
        actual_norm, actual_mean, actual_stdev = \
            torch.ops._caffe2.LayerNorm(torch.tensor(X), torch.tensor(
                weight), torch.tensor(bias), 1, epsilon, True)
        torch.testing.assert_close(expected_norm, actual_norm)
//  [python]
def test_update(preserve_origin):
    array = torch.randn(32, 4)
    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)
    torch.testing.assert_allclose(transformer.mean, torch.mean(array, 0))
    torch.testing.assert_allclose(transformer.variance, torch.var(array, 0))
//  [python]
def test_update(preserve_origin):
    array = torch.randn(32, 4)
    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)
    torch.testing.assert_close(transformer.mean, torch.mean(array, 0))
    torch.testing.assert_close(transformer.variance, torch.var(array, 0))
//  [python]
def test_normalize(preserve_origin):
    array = torch.randn(32, 4)
    new_array = torch.randn(4)
    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)
    transformed_array = transformer(new_array)

    if preserve_origin:
        mean = 0
        scale = torch.sqrt(torch.var(array, 0) + torch.mean(array, 0) ** 2)
    else:
        mean = torch.mean(array, 0)
        scale = torch.sqrt(torch.var(array, 0))
    torch.testing.assert_allclose(transformed_array, (new_array - mean) / scale)
//  [python]
def test_normalize(preserve_origin):
    array = torch.randn(32, 4)
    new_array = torch.randn(4)
    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)
    transformed_array = transformer(new_array)

    if preserve_origin:
        mean = 0
        scale = torch.sqrt(torch.var(array, 0) + torch.mean(array, 0) ** 2)
    else:
        mean = torch.mean(array, 0)
        scale = torch.sqrt(torch.var(array, 0))
    torch.testing.assert_close(transformed_array, (new_array - mean) / scale)
//  [python]
def test_unnormalize():
    array = torch.randn(32, 4)
    new_array = torch.randn(4)
    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)
    transformed_array = transformer(new_array)
    back_array = transformer.inverse(transformed_array)

    torch.testing.assert_allclose(back_array, new_array)
//  [python]
def test_unnormalize():
    array = torch.randn(32, 4)
    new_array = torch.randn(4)
    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)
    transformed_array = transformer(new_array)
    back_array = transformer.inverse(transformed_array)

    torch.testing.assert_close(back_array, new_array)
//  [python]
def validate_transformed_module(
            # To please flake
            self,
            pattern_count_map,
            data_shape,
            prepack_removal=False,
            fuse_clamping_ops=False):
        module_instance = self
        scripted_model = torch.jit.script(module_instance)
        scripted_model.eval()
        input_data = torch.normal(1, 20, size=data_shape)
        ref_result = scripted_model(input_data)
        torch._C._jit_pass_insert_prepacked_ops(scripted_model._c)
        if fuse_clamping_ops or prepack_removal:
            scripted_model._c = torch._C._freeze_module(scripted_model._c)
        if fuse_clamping_ops:
            torch._C._jit_pass_fuse_clamp_w_prepacked_linear_conv(scripted_model._c)
        if (prepack_removal):
            torch._C._jit_pass_fold_prepacking_ops(scripted_model._c)

        buffer = io.BytesIO()
        torch.jit.save(scripted_model, buffer)
        buffer.seek(0)
        deserialized_scripted_model = torch.jit.load(buffer)
        for pattern, v in pattern_count_map.items():
            if (v == 0):
                FileCheck().check(pattern).run(deserialized_scripted_model.graph)
            elif (v == -1):
                FileCheck().check_not(pattern).run(deserialized_scripted_model.graph)
            else:
                FileCheck().check_count(pattern, v, exactly=True).run(deserialized_scripted_model.graph)
        xnnpack_result = deserialized_scripted_model(input_data)
        torch.testing.assert_allclose(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)
//  [python]
def validate_transformed_module(
            # To please flake
            self,
            pattern_count_map,
            data_shape,
            prepack_removal=False,
            fuse_clamping_ops=False):
        module_instance = self
        scripted_model = torch.jit.script(module_instance)
        scripted_model.eval()
        input_data = torch.normal(1, 20, size=data_shape)
        ref_result = scripted_model(input_data)
        torch._C._jit_pass_insert_prepacked_ops(scripted_model._c)
        if fuse_clamping_ops or prepack_removal:
            scripted_model._c = torch._C._freeze_module(scripted_model._c)
        if fuse_clamping_ops:
            torch._C._jit_pass_fuse_clamp_w_prepacked_linear_conv(scripted_model._c)
        if (prepack_removal):
            torch._C._jit_pass_fold_prepacking_ops(scripted_model._c)

        buffer = io.BytesIO()
        torch.jit.save(scripted_model, buffer)
        buffer.seek(0)
        deserialized_scripted_model = torch.jit.load(buffer)
        for pattern, v in pattern_count_map.items():
            if (v == 0):
                FileCheck().check(pattern).run(deserialized_scripted_model.graph)
            elif (v == -1):
                FileCheck().check_not(pattern).run(deserialized_scripted_model.graph)
            else:
                FileCheck().check_count(pattern, v, exactly=True).run(deserialized_scripted_model.graph)
        xnnpack_result = deserialized_scripted_model(input_data)
        torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)
//  [python]
def _test_rowwise_prune_op(self, embedding_rows, embedding_dims, indices_type, weights_dtype):
        embedding_weights = None
        if weights_dtype in [torch.int8, torch.int16, torch.int32, torch.int64]:
            embedding_weights = torch.randint(0, 100, (embedding_rows, embedding_dims), dtype=weights_dtype)
        else:
            embedding_weights = torch.rand((embedding_rows, embedding_dims), dtype=weights_dtype)
        mask = self._generate_rowwise_mask(embedding_rows)

        def get_pt_result(embedding_weights, mask, indices_type):
            return torch._rowwise_prune(embedding_weights, mask, indices_type)

        # Reference implementation.
        def get_reference_result(embedding_weights, mask, indices_type):
            num_embeddings = mask.size()[0]
            compressed_idx_out = torch.zeros(num_embeddings, dtype=indices_type)
            pruned_weights_out = embedding_weights[mask[:]]
            idx = 0
            for i in range(mask.size()[0]):
                if mask[i]:
                    compressed_idx_out[i] = idx
                    idx = idx + 1
                else:
                    compressed_idx_out[i] = -1
            return (pruned_weights_out, compressed_idx_out)

        pt_pruned_weights, pt_compressed_indices_map = get_pt_result(
            embedding_weights, mask, indices_type)
        ref_pruned_weights, ref_compressed_indices_map = get_reference_result(
            embedding_weights, mask, indices_type)

        torch.testing.assert_allclose(pt_pruned_weights, ref_pruned_weights)
        self.assertEqual(pt_compressed_indices_map, ref_compressed_indices_map)
        self.assertEqual(pt_compressed_indices_map.dtype, indices_type)
//  [python]
def _test_rowwise_prune_op(self, embedding_rows, embedding_dims, indices_type, weights_dtype):
        embedding_weights = None
        if weights_dtype in [torch.int8, torch.int16, torch.int32, torch.int64]:
            embedding_weights = torch.randint(0, 100, (embedding_rows, embedding_dims), dtype=weights_dtype)
        else:
            embedding_weights = torch.rand((embedding_rows, embedding_dims), dtype=weights_dtype)
        mask = self._generate_rowwise_mask(embedding_rows)

        def get_pt_result(embedding_weights, mask, indices_type):
            return torch._rowwise_prune(embedding_weights, mask, indices_type)

        # Reference implementation.
        def get_reference_result(embedding_weights, mask, indices_type):
            num_embeddings = mask.size()[0]
            compressed_idx_out = torch.zeros(num_embeddings, dtype=indices_type)
            pruned_weights_out = embedding_weights[mask[:]]
            idx = 0
            for i in range(mask.size()[0]):
                if mask[i]:
                    compressed_idx_out[i] = idx
                    idx = idx + 1
                else:
                    compressed_idx_out[i] = -1
            return (pruned_weights_out, compressed_idx_out)

        pt_pruned_weights, pt_compressed_indices_map = get_pt_result(
            embedding_weights, mask, indices_type)
        ref_pruned_weights, ref_compressed_indices_map = get_reference_result(
            embedding_weights, mask, indices_type)

        torch.testing.assert_close(pt_pruned_weights, ref_pruned_weights)
        self.assertEqual(pt_compressed_indices_map, ref_compressed_indices_map)
        self.assertEqual(pt_compressed_indices_map.dtype, indices_type)
//  [python]
def test_sequential_update():
    torch.manual_seed(0)
    array = torch.randn(16, 4)

    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)

    torch.testing.assert_allclose(transformer.mean, torch.mean(array, 0))
    torch.testing.assert_allclose(transformer.variance, torch.var(array, 0))

    for _ in range(10):
        new_array = torch.randn(torch.randint(1, 32, (1,)), 4)
        transformer.update(new_array)

        array = torch.cat((array, new_array), dim=0)
        torch.testing.assert_allclose(transformer.mean, torch.mean(array, 0))
        torch.testing.assert_allclose(transformer.variance, torch.var(array, 0))
//  [python]
def test_sequential_update():
    torch.manual_seed(0)
    array = torch.randn(16, 4)

    transformer = Normalizer(dim=(4,), preserve_origin=preserve_origin)
    transformer.update(array)

    torch.testing.assert_close(transformer.mean, torch.mean(array, 0))
    torch.testing.assert_close(transformer.variance, torch.var(array, 0))

    for _ in range(10):
        new_array = torch.randn(torch.randint(1, 32, (1,)), 4)
        transformer.update(new_array)

        array = torch.cat((array, new_array), dim=0)
        torch.testing.assert_close(transformer.mean, torch.mean(array, 0))
        torch.testing.assert_close(transformer.variance, torch.var(array, 0))
//  [python]
def validate_transform_conv1d_to_conv2d(
            self,
            pattern_count_transformed_map,
            pattern_count_optimized_map,
            data_shape):
        module_instance = self
        scripted_model = torch.jit.script(module_instance)
        scripted_model.eval()
        input_data = torch.normal(1, 20, size=data_shape)
        ref_result = scripted_model(input_data)
        torch._C._jit_pass_transform_conv1d_to_conv2d(scripted_model._c)
        optimized_scripted_model = optimize_for_mobile(scripted_model)

        buffer = io.BytesIO()
        torch.jit.save(scripted_model, buffer)
        buffer.seek(0)
        deserialized_scripted_model = torch.jit.load(buffer)

        for pattern, v in pattern_count_transformed_map.items():
            if (v == 0):
                FileCheck().check(pattern).run(deserialized_scripted_model.graph)
            elif (v == -1):
                FileCheck().check_not(pattern).run(deserialized_scripted_model.graph)
            else:
                FileCheck().check_count(pattern, v, exactly=True).run(deserialized_scripted_model.graph)
        transformed_result = deserialized_scripted_model(input_data)
        torch.testing.assert_allclose(ref_result, transformed_result, rtol=1e-2, atol=1e-3)

        optimized_buffer = io.BytesIO()
        torch.jit.save(optimized_scripted_model, optimized_buffer)
        optimized_buffer.seek(0)
        deserialized_optimized_scripted_model = torch.jit.load(optimized_buffer)

        for pattern, v in pattern_count_optimized_map.items():
            if (v == 0):
                FileCheck().check(pattern).run(deserialized_optimized_scripted_model.graph)
            elif (v == -1):
                FileCheck().check_not(pattern).run(deserialized_optimized_scripted_model.graph)
            else:
                FileCheck().check_count(pattern, v, exactly=True).run(deserialized_optimized_scripted_model.graph)
        xnnpack_result = deserialized_optimized_scripted_model(input_data)
        torch.testing.assert_allclose(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)
//  [python]
def validate_transform_conv1d_to_conv2d(
            self,
            pattern_count_transformed_map,
            pattern_count_optimized_map,
            data_shape):
        module_instance = self
        scripted_model = torch.jit.script(module_instance)
        scripted_model.eval()
        input_data = torch.normal(1, 20, size=data_shape)
        ref_result = scripted_model(input_data)
        torch._C._jit_pass_transform_conv1d_to_conv2d(scripted_model._c)
        optimized_scripted_model = optimize_for_mobile(scripted_model)

        buffer = io.BytesIO()
        torch.jit.save(scripted_model, buffer)
        buffer.seek(0)
        deserialized_scripted_model = torch.jit.load(buffer)

        for pattern, v in pattern_count_transformed_map.items():
            if (v == 0):
                FileCheck().check(pattern).run(deserialized_scripted_model.graph)
            elif (v == -1):
                FileCheck().check_not(pattern).run(deserialized_scripted_model.graph)
            else:
                FileCheck().check_count(pattern, v, exactly=True).run(deserialized_scripted_model.graph)
        transformed_result = deserialized_scripted_model(input_data)
        torch.testing.assert_close(ref_result, transformed_result, rtol=1e-2, atol=1e-3)

        optimized_buffer = io.BytesIO()
        torch.jit.save(optimized_scripted_model, optimized_buffer)
        optimized_buffer.seek(0)
        deserialized_optimized_scripted_model = torch.jit.load(optimized_buffer)

        for pattern, v in pattern_count_optimized_map.items():
            if (v == 0):
                FileCheck().check(pattern).run(deserialized_optimized_scripted_model.graph)
            elif (v == -1):
                FileCheck().check_not(pattern).run(deserialized_optimized_scripted_model.graph)
            else:
                FileCheck().check_count(pattern, v, exactly=True).run(deserialized_optimized_scripted_model.graph)
        xnnpack_result = deserialized_optimized_scripted_model(input_data)
        torch.testing.assert_close(ref_result, xnnpack_result, rtol=1e-2, atol=1e-3)
//  [python]
def check_parity(torch_bn, fs_bn, x):
    yh = torch.randn_like(x)
    torch_x = x.detach()
    torch_x.requires_grad = True
    torch_y = torch_bn(torch_x)
    torch_y.backward(yh)
    fs_x = x.detach()
    fs_x.requires_grad = True
    fs_y = fs_bn(fs_x)
    fs_y.backward(yh)
    torch.testing.assert_allclose(torch_y, fs_y)
    torch.testing.assert_allclose(torch_bn.running_mean, fs_bn.running_mean)
    torch.testing.assert_allclose(torch_bn.running_var, fs_bn.running_var)
    torch.testing.assert_allclose(torch_bn.weight, fs_bn.weight)
    torch.testing.assert_allclose(torch_bn.bias, fs_bn.bias)
    torch.testing.assert_allclose(torch_bn.weight.grad, fs_bn.weight.grad)
    torch.testing.assert_allclose(torch_bn.bias.grad, fs_bn.bias.grad)
    torch.testing.assert_allclose(torch_x.grad, fs_x.grad)
//  [python]
def check_parity(torch_bn, fs_bn, x):
    yh = torch.randn_like(x)
    torch_x = x.detach()
    torch_x.requires_grad = True
    torch_y = torch_bn(torch_x)
    torch_y.backward(yh)
    fs_x = x.detach()
    fs_x.requires_grad = True
    fs_y = fs_bn(fs_x)
    fs_y.backward(yh)
    torch.testing.assert_close(torch_y, fs_y)
    torch.testing.assert_close(torch_bn.running_mean, fs_bn.running_mean)
    torch.testing.assert_close(torch_bn.running_var, fs_bn.running_var)
    torch.testing.assert_close(torch_bn.weight, fs_bn.weight)
    torch.testing.assert_close(torch_bn.bias, fs_bn.bias)
    torch.testing.assert_close(torch_bn.weight.grad, fs_bn.weight.grad)
    torch.testing.assert_close(torch_bn.bias.grad, fs_bn.bias.grad)
    torch.testing.assert_close(torch_x.grad, fs_x.grad)
//  [python]
def test_ragged():
    vals = torch.as_tensor([0, 1, 2, 3, 4, 5]).float()
    widths = torch.as_tensor([3, 1, 2]).int()
    ragged = Ragged(vals, widths)

    torch.testing.assert_allclose(ragged[1], [3])
    torch.testing.assert_allclose(ragged[-1], [4, 5])

    torch.testing.assert_allclose(ragged[:2].vals, [0, 1, 2, 3])
    torch.testing.assert_allclose(ragged[:2].widths, [3, 1])

    torch.testing.assert_allclose(ragged[1:].vals, np.array([3, 4, 5]))
    torch.testing.assert_allclose(ragged[1:].widths, np.array([1, 2]))
//  [python]
def test_ragged():
    vals = torch.as_tensor([0, 1, 2, 3, 4, 5]).float()
    widths = torch.as_tensor([3, 1, 2]).int()
    ragged = Ragged(vals, widths)

    torch.testing.assert_close(ragged[1], [3])
    torch.testing.assert_close(ragged[-1], [4, 5])

    torch.testing.assert_close(ragged[:2].vals, [0, 1, 2, 3])
    torch.testing.assert_close(ragged[:2].widths, [3, 1])

    torch.testing.assert_close(ragged[1:].vals, np.array([3, 4, 5]))
    torch.testing.assert_close(ragged[1:].widths, np.array([1, 2]))
//  [python]
def test_interp():
    n = 10
    c = 3
    d = 4
    t = 5
    x = torch.linspace(0, 1, steps=n)
    x = x[None, :].repeat((c, 1))
    y = torch.randn(c, d, n)
    xs = x[:, t].reshape(c, 1, 1)
    ind = torch.tensor(t, dtype=torch.long).reshape(1, 1, 1).repeat((c, 1, 1))
    ys = cubicspline.interp(x, y, xs, ind)
    yy = torch.stack([y[i, :, ind[i]] for i in range(c)], dim=0)
    torch.testing.assert_allclose(yy, ys)
//  [python]
def test_interp():
    n = 10
    c = 3
    d = 4
    t = 5
    x = torch.linspace(0, 1, steps=n)
    x = x[None, :].repeat((c, 1))
    y = torch.randn(c, d, n)
    xs = x[:, t].reshape(c, 1, 1)
    ind = torch.tensor(t, dtype=torch.long).reshape(1, 1, 1).repeat((c, 1, 1))
    ys = cubicspline.interp(x, y, xs, ind)
    yy = torch.stack([y[i, :, ind[i]] for i in range(c)], dim=0)
    torch.testing.assert_close(yy, ys)
//  [python]
def check_parity_ddp(torch_bn, fs_bn, x):
    yh = torch.randn_like(x)
    rank = dist.get_rank()
    torch_ddp = DDP(torch_bn, device_ids=[rank])
    torch_bn = torch_ddp.module
    torch_x = x.detach()
    torch_x.requires_grad = True
    torch_y = torch_ddp(torch_x)
    torch_y.backward(yh)
    fs_ddp = DDP(fs_bn, device_ids=[rank])
    fs_bn = fs_ddp.module
    fs_x = x.detach()
    fs_x.requires_grad = True
    fs_y = fs_ddp(fs_x)
    fs_y.backward(yh)
    torch.testing.assert_allclose(torch_y, fs_y)
    torch.testing.assert_allclose(torch_x.grad, fs_x.grad)
    if isinstance(torch_bn, nn.Sequential):
        torch_bn = torch_bn[0]
        fs_bn = fs_bn[0]
    torch.testing.assert_allclose(torch_bn.running_mean, fs_bn.running_mean)
    torch.testing.assert_allclose(torch_bn.running_var, fs_bn.running_var)
    torch.testing.assert_allclose(torch_bn.weight, fs_bn.weight)
    torch.testing.assert_allclose(torch_bn.bias, fs_bn.bias)
    torch.testing.assert_allclose(torch_bn.weight.grad, fs_bn.weight.grad)
    torch.testing.assert_allclose(torch_bn.bias.grad, fs_bn.bias.grad)