jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_contiguous_and_stride_order(executor: TestExecutor, device: str, _):
    inp = torch.randn(2, 4, 5, 3, device=device, dtype=torch.float32).permute(0, 3, 1, 2)

    def foo(a, order):
        return clang.stride_order(a, order)

    # stride order, expected strides (from shape 2, 4, 5, 3)
    test_cases = (
        ((3, 2, 1, 0), (60, 20, 5, 1)),
        ((3, 0, 2, 1), (60, 1, 15, 3)),
    )

    for stride_order, expected_strides in test_cases:
        cfoo = executor.make_callable(foo)
        o = cfoo(inp, stride_order)
        assert o.shape == inp.shape
        assert o.stride() == expected_strides
        t.assert_close(inp, o)

    # order=none case
    thunder_result = cfoo(inp, None)
    assert thunder_result.stride() == (60, 20, 5, 1)
    t.assert_close(thunder_result, inp)

    # Memory format tests
    def channels_last_2d(a, memory_format):
        return a.contiguous(memory_format=memory_format)

    cfn = executor.make_callable(channels_last_2d, disable_preprocessing=False)

    # Contiguous cases
    a = torch.randn((4, 3, 2), device=device, dtype=torch.float32)
    thunder_result = cfn(a, torch.contiguous_format)
    torch_result = channels_last_2d(a, torch.contiguous_format)
    t.assert_close(torch_result, thunder_result, check_stride=True)

    b = a.permute(0, 2, 1)
    thunder_result = cfn(a, torch.contiguous_format)
    torch_result = channels_last_2d(a, torch.contiguous_format)
    t.assert_close(torch_result, thunder_result, check_stride=True)

    # Channels last 2D cases
    a = torch.randn((5, 4, 3, 2), device=device, dtype=torch.float32)
    thunder_result = cfn(a, torch.channels_last)
    torch_result = channels_last_2d(a, torch.channels_last)
    t.assert_close(torch_result, thunder_result, check_stride=True)

    b = a.permute(3, 1, 0, 2)
    thunder_result = cfn(b, torch.channels_last)
    torch_result = channels_last_2d(b, torch.channels_last)
    t.assert_close(torch_result, thunder_result, check_stride=True)

    # Channels last 3D cases
    a = torch.randn((5, 4, 3, 7, 2), device=device, dtype=torch.float32)
    thunder_result = cfn(a, torch.channels_last_3d)
    torch_result = channels_last_2d(a, torch.channels_last_3d)
    t.assert_close(torch_result, thunder_result, check_stride=True)

    b = a.permute(0, 4, 2, 1, 3)
    thunder_result = cfn(a, torch.channels_last_3d)
    torch_result = channels_last_2d(a, torch.channels_last_3d)
    t.assert_close(torch_result, thunder_result, check_stride=True)
//  [python]
def test_contiguous_and_stride_order(executor: TestExecutor, device: str, _):
    inp = torch.randn(2, 4, 5, 3, device=device, dtype=torch.float32).permute(0, 3, 1, 2)

    def foo(a, order):
        return clang.stride_order(a, order)

    # stride order, expected strides (from shape 2, 4, 5, 3)
    test_cases = (
        ((3, 2, 1, 0), (60, 20, 5, 1)),
        ((3, 0, 2, 1), (60, 1, 15, 3)),
    )

    for stride_order, expected_strides in test_cases:
        cfoo = executor.make_callable(foo)
        o = cfoo(inp, stride_order)
        assert o.shape == inp.shape
        assert o.stride() == expected_strides
        assert_close(inp, o)

    # order=none case
    thunder_result = cfoo(inp, None)
    assert thunder_result.stride() == (60, 20, 5, 1)
    assert_close(thunder_result, inp)

    # Memory format tests
    def channels_last_2d(a, memory_format):
        return a.contiguous(memory_format=memory_format)

    cfn = executor.make_callable(channels_last_2d, disable_preprocessing=False)

    # Contiguous cases
    a = torch.randn((4, 3, 2), device=device, dtype=torch.float32)
    thunder_result = cfn(a, torch.contiguous_format)
    torch_result = channels_last_2d(a, torch.contiguous_format)
    assert_close(torch_result, thunder_result, check_stride=True)

    b = a.permute(0, 2, 1)
    thunder_result = cfn(a, torch.contiguous_format)
    torch_result = channels_last_2d(a, torch.contiguous_format)
    assert_close(torch_result, thunder_result, check_stride=True)

    # Channels last 2D cases
    a = torch.randn((5, 4, 3, 2), device=device, dtype=torch.float32)
    thunder_result = cfn(a, torch.channels_last)
    torch_result = channels_last_2d(a, torch.channels_last)
    assert_close(torch_result, thunder_result, check_stride=True)

    b = a.permute(3, 1, 0, 2)
    thunder_result = cfn(b, torch.channels_last)
    torch_result = channels_last_2d(b, torch.channels_last)
    assert_close(torch_result, thunder_result, check_stride=True)

    # Channels last 3D cases
    a = torch.randn((5, 4, 3, 7, 2), device=device, dtype=torch.float32)
    thunder_result = cfn(a, torch.channels_last_3d)
    torch_result = channels_last_2d(a, torch.channels_last_3d)
    assert_close(torch_result, thunder_result, check_stride=True)

    b = a.permute(0, 4, 2, 1, 3)
    thunder_result = cfn(a, torch.channels_last_3d)
    torch_result = channels_last_2d(a, torch.channels_last_3d)
    assert_close(torch_result, thunder_result, check_stride=True)
//  [python]
def test_incorporate_extra_connections_into_inter_res_conn_set(torch_device):
    n_poses, max_n_blocks, max_n_conn = 2, 5, 3
    resolved_expoly_connections = [[(2, 2, 4, 2)], [(1, 2, 3, 2)]]
    inter_residue_connections64 = torch.full(
        (n_poses, max_n_blocks, max_n_conn, 2),
        -1,
        dtype=torch.int64,
        device=torch_device,
    )
    inter_residue_connections64_gold = inter_residue_connections64.clone()

    PoseStackBuilder._incorporate_extra_connections_into_inter_res_conn_set(
        resolved_expoly_connections, inter_residue_connections64
    )

    inter_residue_connections64_gold[0, 2, 2, 0] = 4
    inter_residue_connections64_gold[0, 2, 2, 1] = 2
    inter_residue_connections64_gold[0, 4, 2, 0] = 2
    inter_residue_connections64_gold[0, 4, 2, 1] = 2

    inter_residue_connections64_gold[1, 1, 2, 0] = 3
    inter_residue_connections64_gold[1, 1, 2, 1] = 2
    inter_residue_connections64_gold[1, 3, 2, 0] = 1
    inter_residue_connections64_gold[1, 3, 2, 1] = 2

    t.assert_allclose(
        inter_residue_connections64_gold, inter_residue_connections64
    )
//  [python]
def test_incorporate_extra_connections_into_inter_res_conn_set(torch_device):
    n_poses, max_n_blocks, max_n_conn = 2, 5, 3
    resolved_expoly_connections = [[(2, 2, 4, 2)], [(1, 2, 3, 2)]]
    inter_residue_connections64 = torch.full(
        (n_poses, max_n_blocks, max_n_conn, 2),
        -1,
        dtype=torch.int64,
        device=torch_device,
    )
    inter_residue_connections64_gold = inter_residue_connections64.clone()

    PoseStackBuilder._incorporate_extra_connections_into_inter_res_conn_set(
        resolved_expoly_connections, inter_residue_connections64
    )

    inter_residue_connections64_gold[0, 2, 2, 0] = 4
    inter_residue_connections64_gold[0, 2, 2, 1] = 2
    inter_residue_connections64_gold[0, 4, 2, 0] = 2
    inter_residue_connections64_gold[0, 4, 2, 1] = 2

    inter_residue_connections64_gold[1, 1, 2, 0] = 3
    inter_residue_connections64_gold[1, 1, 2, 1] = 2
    inter_residue_connections64_gold[1, 3, 2, 0] = 1
    inter_residue_connections64_gold[1, 3, 2, 1] = 2

    torch.testing.assert_close(
        inter_residue_connections64_gold, inter_residue_connections64
    )
//  [python]
def test_compute_pitch_batched(self, test_data_dir):
        audio_filepath = Path(test_data_dir) / "tts" / "mini_ljspeech" / "wavs" / "LJ003-0182.wav"
        manifest_entry = {"audio_filepath": audio_filepath}

        # Compute pitch with batching disabled
        pitch_featurizer = PitchFeaturizer(hop_length=self.hop_len, sample_rate=self.sample_rate, batch_seconds=None)
        pitch, voiced, voiced_prob = pitch_featurizer.compute_pitch(
            manifest_entry=manifest_entry, audio_dir=test_data_dir
        )

        # Compute pitch in 1 second chunks
        pitch_featurizer_batch = PitchFeaturizer(
            hop_length=self.hop_len, sample_rate=self.sample_rate, batch_seconds=1.0
        )
        pitch_batch, voiced_batch, voiced_prob_batch = pitch_featurizer_batch.compute_pitch(
            manifest_entry=manifest_entry, audio_dir=test_data_dir
        )

        t.assert_allclose(pitch_batch, pitch)
        t.assert_allclose(voiced_batch, voiced)
        t.assert_allclose(voiced_prob_batch, voiced_prob)
//  [python]
def test_compute_pitch_batched(self, test_data_dir):
        audio_filepath = Path(test_data_dir) / "tts" / "mini_ljspeech" / "wavs" / "LJ003-0182.wav"
        manifest_entry = {"audio_filepath": audio_filepath}

        # Compute pitch with batching disabled
        pitch_featurizer = PitchFeaturizer(hop_length=self.hop_len, sample_rate=self.sample_rate, batch_seconds=None)
        pitch, voiced, voiced_prob = pitch_featurizer.compute_pitch(
            manifest_entry=manifest_entry, audio_dir=test_data_dir
        )

        # Compute pitch in 1 second chunks
        pitch_featurizer_batch = PitchFeaturizer(
            hop_length=self.hop_len, sample_rate=self.sample_rate, batch_seconds=1.0
        )
        pitch_batch, voiced_batch, voiced_prob_batch = pitch_featurizer_batch.compute_pitch(
            manifest_entry=manifest_entry, audio_dir=test_data_dir
        )

        torch.testing.assert_close(pitch_batch, pitch)
        torch.testing.assert_close(voiced_batch, voiced)
        torch.testing.assert_close(voiced_prob_batch, voiced_prob)
//  [python]
def _run_test(self, dtypes) -> None:
        (
            indices_list,
            offsets_list,
            per_sample_weights,
            empty_per_sample_weights,
            include_last_offsets,
        ) = self._get_inputs(dtypes)
        ref_mod = TBEInputPrepareReference(include_last_offsets)

        outputs = torch.ops.fbgemm.tbe_input_combine(
            indices_list,
            offsets_list,
            per_sample_weights,
            torch.BoolTensor(include_last_offsets),
        )
        ref_outputs = ref_mod(indices_list, offsets_list, per_sample_weights)
        for i, j in zip(outputs, ref_outputs):
            self.assertEqual(i,j)
        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)

        outputs = torch.ops.fbgemm.tbe_input_combine(
            indices_list,
            offsets_list,
            empty_per_sample_weights,
            torch.BoolTensor(include_last_offsets),
        )
        ref_outputs = ref_mod(indices_list, offsets_list, empty_per_sample_weights)
        for i, j in zip(outputs[:-1], ref_outputs[:-1]):
            self.assertEqual(i,j)
            self.assertTrue(j.dtype == torch.int32)

        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)
        self.assertTrue(outputs[-1].size(0) == 0)
//  [python]
def _run_test(self, dtypes) -> None:
        (
            indices_list,
            offsets_list,
            per_sample_weights,
            empty_per_sample_weights,
            include_last_offsets,
        ) = self._get_inputs(dtypes)
        ref_mod = TBEInputPrepareReference(include_last_offsets)

        outputs = torch.ops.fbgemm.tbe_input_combine(
            indices_list,
            offsets_list,
            per_sample_weights,
            torch.BoolTensor(include_last_offsets),
        )
        ref_outputs = ref_mod(indices_list, offsets_list, per_sample_weights)
        for i, j in zip(outputs, ref_outputs):
            torch.testing.assert_close(i, j)
        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)

        outputs = torch.ops.fbgemm.tbe_input_combine(
            indices_list,
            offsets_list,
            empty_per_sample_weights,
            torch.BoolTensor(include_last_offsets),
        )
        ref_outputs = ref_mod(indices_list, offsets_list, empty_per_sample_weights)
        for i, j in zip(outputs[:-1], ref_outputs[:-1]):
            torch.testing.assert_close(i, j)
            self.assertTrue(j.dtype == torch.int32)

        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)
        self.assertTrue(outputs[-1].size(0) == 0)
//  [python]
def _run_padding_fused_test(self, dtypes, batch_size) -> None:
        (
            indices_list,
            offsets_list,
            per_sample_weights,
            empty_per_sample_weights,
            include_last_offsets,
        ) = self._get_inputs(dtypes)
        ref_mod = TBEInputPrepareReference(include_last_offsets)

        outputs = torch.ops.fbgemm.padding_fused_tbe_input_combine(
            indices_list,
            offsets_list,
            per_sample_weights,
            torch.BoolTensor(include_last_offsets),
            batch_size,
        )
        ref_outputs = ref_mod(
            indices_list, offsets_list, per_sample_weights, batch_size
        )
        for i, j in zip(outputs, ref_outputs):
            torch.testing.assert_allclose(i, j)
        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)

        outputs = torch.ops.fbgemm.padding_fused_tbe_input_combine(
            indices_list,
            offsets_list,
            empty_per_sample_weights,
            torch.BoolTensor(include_last_offsets),
            batch_size,
        )
        ref_outputs = ref_mod(
            indices_list, offsets_list, empty_per_sample_weights, batch_size
        )
        for i, j in zip(outputs[:-1], ref_outputs[:-1]):
            torch.testing.assert_allclose(i, j)
            self.assertTrue(j.dtype == torch.int32)

        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)
        self.assertTrue(outputs[-1].size(0) == 0)
//  [python]
def _run_padding_fused_test(self, dtypes, batch_size) -> None:
        (
            indices_list,
            offsets_list,
            per_sample_weights,
            empty_per_sample_weights,
            include_last_offsets,
        ) = self._get_inputs(dtypes)
        ref_mod = TBEInputPrepareReference(include_last_offsets)

        outputs = torch.ops.fbgemm.padding_fused_tbe_input_combine(
            indices_list,
            offsets_list,
            per_sample_weights,
            torch.BoolTensor(include_last_offsets),
            batch_size,
        )
        ref_outputs = ref_mod(
            indices_list, offsets_list, per_sample_weights, batch_size
        )
        for i, j in zip(outputs, ref_outputs):
            torch.testing.assert_close(i, j)
        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)

        outputs = torch.ops.fbgemm.padding_fused_tbe_input_combine(
            indices_list,
            offsets_list,
            empty_per_sample_weights,
            torch.BoolTensor(include_last_offsets),
            batch_size,
        )
        ref_outputs = ref_mod(
            indices_list, offsets_list, empty_per_sample_weights, batch_size
        )
        for i, j in zip(outputs[:-1], ref_outputs[:-1]):
            torch.testing.assert_close(i, j)
            self.assertTrue(j.dtype == torch.int32)

        self.assertTrue(outputs[0].dtype == torch.int32)
        self.assertTrue(outputs[1].dtype == torch.int32)
        self.assertTrue(outputs[-1].size(0) == 0)
//  [python]
def test_gradient_promoting_in_fp32(half_precision: torch.dtype):
    model = nn.Linear(3, 2, bias=False, dtype=half_precision, device="cuda")

    # Create Nanotron Parameter
    model.weight = NanotronParameter(model.weight)

    # Add gradient accumulator
    accumulator = FP32GradientAccumulator(model.named_parameters())

    # Check that our model is a valid model
    sanity_check(model)

    # Compute backward
    input = torch.randn(5, 3, dtype=half_precision, device="cuda")
    accumulator.backward(model(input).sum())

    # Check that we have an high precision gradient and that the low precision one is cleared
    assert accumulator.parameters["weight"]["fp32"].grad.dtype == torch.float
    if model.weight.grad is not None:
        # We check that it's zero
        torch.testing.assert_allclose(model.weight.grad, torch.zeros_like(model.weight.grad), atol=1e-6, rtol=1e-7)
//  [python]
def test_gradient_promoting_in_fp32(half_precision: torch.dtype):
    model = nn.Linear(3, 2, bias=False, dtype=half_precision, device="cuda")

    # Create Nanotron Parameter
    model.weight = NanotronParameter(model.weight)

    # Add gradient accumulator
    accumulator = FP32GradientAccumulator(model.named_parameters())

    # Check that our model is a valid model
    sanity_check(model)

    # Compute backward
    input = torch.randn(5, 3, dtype=half_precision, device="cuda")
    accumulator.backward(model(input).sum())

    # Check that we have an high precision gradient and that the low precision one is cleared
    assert accumulator.parameters["weight"]["fp32"].grad.dtype == torch.float
    if model.weight.grad is not None:
        # We check that it's zero
        torch.testing.assert_close(model.weight.grad, torch.zeros_like(model.weight.grad), atol=1e-6, rtol=1e-7)
//  [python]
def test_npu_adamw_compare_with_cpu(self):
        from atorch.npu.optim import NpuAdamW

        weight_value = torch.randn(10, 5)
        bias_value = torch.randn(10)
        input_value = torch.randn(5)
        weight_cpu = Variable(weight_value, requires_grad=True)
        bias_cpu = Variable(bias_value, requires_grad=True)
        input_cpu = Variable(input_value)

        weight_npu = Variable(weight_value.npu(), requires_grad=True)
        bias_npu = Variable(bias_value.npu(), requires_grad=True)
        input_npu = Variable(input_value.npu(), requires_grad=True)

        ori_optim = torch.optim.AdamW([weight_cpu, bias_cpu], lr=1e-3)
        ori_optim.zero_grad()
        y_cpu = weight_cpu.mv(input_cpu)
        loss_cpu = (y_cpu + bias_cpu).pow(2).sum()
        loss_cpu.backward()
        ori_optim.step()

        npu_optim = NpuAdamW([weight_npu, bias_npu], lr=1e-3)
        npu_optim.zero_grad()
        y_npu = weight_npu.mv(input_npu)
        loss_npu = (y_npu + bias_npu).pow(2).sum()
        loss_npu.backward()
        npu_optim.step()

        torch.testing.assert_allclose(weight_npu.cpu(), weight_cpu)
        torch.testing.assert_allclose(bias_npu.cpu(), bias_cpu)
//  [python]
def test_npu_adamw_compare_with_cpu(self):
        from atorch.npu.optim import NpuAdamW

        weight_value = torch.randn(10, 5)
        bias_value = torch.randn(10)
        input_value = torch.randn(5)
        weight_cpu = Variable(weight_value, requires_grad=True)
        bias_cpu = Variable(bias_value, requires_grad=True)
        input_cpu = Variable(input_value)

        weight_npu = Variable(weight_value.npu(), requires_grad=True)
        bias_npu = Variable(bias_value.npu(), requires_grad=True)
        input_npu = Variable(input_value.npu(), requires_grad=True)

        ori_optim = torch.optim.AdamW([weight_cpu, bias_cpu], lr=1e-3)
        ori_optim.zero_grad()
        y_cpu = weight_cpu.mv(input_cpu)
        loss_cpu = (y_cpu + bias_cpu).pow(2).sum()
        loss_cpu.backward()
        ori_optim.step()

        npu_optim = NpuAdamW([weight_npu, bias_npu], lr=1e-3)
        npu_optim.zero_grad()
        y_npu = weight_npu.mv(input_npu)
        loss_npu = (y_npu + bias_npu).pow(2).sum()
        loss_npu.backward()
        npu_optim.step()

        torch.testing.assert_close(weight_npu.cpu(), weight_cpu)
        torch.testing.assert_close(bias_npu.cpu(), bias_cpu)
//  [python]
def test_batch_der(target, prediction):
    der = diarization_error_rate(prediction, target, reduce="batch")
    t.assert_close(der.item(), (3.0 + 2.0 + 1.0) / 12.0)
//  [python]
def test_batch_der(target, prediction):
    der = diarization_error_rate(prediction, target, reduce="batch")
    torch.testing.assert_close(der.item(), (3.0 + 2.0 + 1.0) / 12.0)
//  [python]
def test_gradient_accumulated_in_fp32(half_precision: torch.dtype):
    model = nn.Linear(3, 2, bias=False, dtype=half_precision, device="cuda")
    ref_model = nn.Linear(3, 2, bias=False, dtype=half_precision, device="cuda")
    with torch.inference_mode():
        ref_model.weight.copy_(model.weight)

    # Create Nanotron Parameter
    model.weight = NanotronParameter(model.weight)

    # Add gradient accumulator
    accumulator = FP32GradientAccumulator(model.named_parameters())

    # Check that our model is a valid model
    sanity_check(model)

    # Compute backward
    grad_accumulation_steps = 2
    for _ in range(grad_accumulation_steps):
        # We want large input to have large gradients.
        input = (torch.randn(5, 3, dtype=half_precision, device="cuda") ** 2 + 1) * 100

        # Compute backwards
        accumulator.backward(model(input).sum())
        ref_model(input).sum().backward()

    # We check that we get the same gradient accumulation. In theory we do get more precision by promoting gradients to fp32.
    torch.testing.assert_allclose(
        accumulator.parameters["weight"]["fp32"].grad.to(half_precision),
        ref_model.weight.grad,
    )
//  [python]
def test_gradient_accumulated_in_fp32(half_precision: torch.dtype):
    model = nn.Linear(3, 2, bias=False, dtype=half_precision, device="cuda")
    ref_model = nn.Linear(3, 2, bias=False, dtype=half_precision, device="cuda")
    with torch.inference_mode():
        ref_model.weight.copy_(model.weight)

    # Create Nanotron Parameter
    model.weight = NanotronParameter(model.weight)

    # Add gradient accumulator
    accumulator = FP32GradientAccumulator(model.named_parameters())

    # Check that our model is a valid model
    sanity_check(model)

    # Compute backward
    grad_accumulation_steps = 2
    for _ in range(grad_accumulation_steps):
        # We want large input to have large gradients.
        input = (torch.randn(5, 3, dtype=half_precision, device="cuda") ** 2 + 1) * 100

        # Compute backwards
        accumulator.backward(model(input).sum())
        ref_model(input).sum().backward()

    # We check that we get the same gradient accumulation. In theory we do get more precision by promoting gradients to fp32.
    torch.testing.assert_close(
        accumulator.parameters["weight"]["fp32"].grad.to(half_precision),
        ref_model.weight.grad,
    )
//  [python]
def test_batch_reduction(target, prediction):
    false_alarm, missed_detection, speaker_confusion, speech_total = _der_update(
        prediction, target, reduce="batch"
    )
    t.assert_close(false_alarm.item(), 3.0)
    t.assert_close(missed_detection.item(), 2.0)
    t.assert_close(speaker_confusion.item(), 1.0)
    t.assert_close(speech_total.item(), 12.0)
//  [python]
def test_batch_reduction(target, prediction):
    false_alarm, missed_detection, speaker_confusion, speech_total = _der_update(
        prediction, target, reduce="batch"
    )
    torch.testing.assert_close(false_alarm.item(), 3.0)
    torch.testing.assert_close(missed_detection.item(), 2.0)
    torch.testing.assert_close(speaker_confusion.item(), 1.0)
    torch.testing.assert_close(speech_total.item(), 12.0)
//  [python]
def test_chunk_reduction(target, prediction):
    false_alarm, missed_detection, speaker_confusion, speech_total = _der_update(
        prediction, target, reduce="chunk"
    )

    t.assert_close(
        false_alarm,
        torch.Tensor([1.0, 2.0]),
    )

    t.assert_close(
        missed_detection,
        torch.Tensor([2.0, 0.0]),
    )

    t.assert_close(
        speaker_confusion,
        torch.Tensor([1.0, 0.0]),
    )

    t.assert_close(
        speech_total,
        torch.Tensor([8.0, 4.0]),
    )
//  [python]
def test_chunk_reduction(target, prediction):
    false_alarm, missed_detection, speaker_confusion, speech_total = _der_update(
        prediction, target, reduce="chunk"
    )

    torch.testing.assert_close(
        false_alarm,
        torch.Tensor([1.0, 2.0]),
    )

    torch.testing.assert_close(
        missed_detection,
        torch.Tensor([2.0, 0.0]),
    )

    torch.testing.assert_close(
        speaker_confusion,
        torch.Tensor([1.0, 0.0]),
    )

    torch.testing.assert_close(
        speech_total,
        torch.Tensor([8.0, 4.0]),
    )
//  [python]
def test_chunk_der(target, prediction):
    der = diarization_error_rate(prediction, target, reduce="chunk")
    t.assert_close(der, torch.Tensor([4.0 / 8.0, 2.0 / 4.0]))
//  [python]
def test_chunk_der(target, prediction):
    der = diarization_error_rate(prediction, target, reduce="chunk")
    torch.testing.assert_close(der, torch.Tensor([4.0 / 8.0, 2.0 / 4.0]))
//  [python]
def test_frame_reduction(target, prediction):
    false_alarm, missed_detection, speaker_confusion, speech_total = _der_update(
        prediction, target, reduce="frame"
    )

    t.assert_close(
        false_alarm,
        torch.Tensor(
            [[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]]
        ),
    )

    t.assert_close(
        missed_detection,
        torch.Tensor(
            [
                [0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            ]
        ),
    )

    t.assert_close(
        speaker_confusion,
        torch.Tensor(
            [[0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]
        ),
    )

    t.assert_close(
        speech_total,
        torch.Tensor(
            [[0.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0]]
        ),
    )
//  [python]
def test_frame_reduction(target, prediction):
    false_alarm, missed_detection, speaker_confusion, speech_total = _der_update(
        prediction, target, reduce="frame"
    )

    torch.testing.assert_close(
        false_alarm,
        torch.Tensor(
            [[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]]
        ),
    )

    torch.testing.assert_close(
        missed_detection,
        torch.Tensor(
            [
                [0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            ]
        ),
    )

    torch.testing.assert_close(
        speaker_confusion,
        torch.Tensor(
            [[0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]
        ),
    )

    torch.testing.assert_close(
        speech_total,
        torch.Tensor(
            [[0.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0]]
        ),
    )
//  [python]
def test_batch_der_with_components(target, prediction):
    der, (
        false_alarm,
        missed_detection,
        speaker_confusion,
        speech_total,
    ) = diarization_error_rate(
        prediction, target, reduce="batch", return_components=True
    )
    t.assert_close(der.item(), (3.0 + 2.0 + 1.0) / 12.0)
    t.assert_close(false_alarm.item(), 3.0)
    t.assert_close(missed_detection.item(), 2.0)
    t.assert_close(speaker_confusion.item(), 1.0)
    t.assert_close(speech_total.item(), 12.0)
//  [python]
def test_batch_der_with_components(target, prediction):
    der, (
        false_alarm,
        missed_detection,
        speaker_confusion,
        speech_total,
    ) = diarization_error_rate(
        prediction, target, reduce="batch", return_components=True
    )
    torch.testing.assert_close(der.item(), (3.0 + 2.0 + 1.0) / 12.0)
    torch.testing.assert_close(false_alarm.item(), 3.0)
    torch.testing.assert_close(missed_detection.item(), 2.0)
    torch.testing.assert_close(speaker_confusion.item(), 1.0)
    torch.testing.assert_close(speech_total.item(), 12.0)
//  [python]
def test_both_tensors_constant(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = 4
        X3 = 2
        Y = ops.where()(X1, X2, X3, dtype=dtype)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_both_tensors_constant")

        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = 4
            x3_pt = 2
            y_pt = torch.where(x1_pt, x2_pt, x3_pt).to(torch.float32)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt}
            model.run_with_tensors(inputs, [y])
            t.assert_allclose(y_pt, y)
//  [python]
def test_both_tensors_constant(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = 4
        X3 = 2
        Y = ops.where()(X1, X2, X3, dtype=dtype)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_both_tensors_constant")

        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = 4
            x3_pt = 2
            y_pt = torch.where(x1_pt, x2_pt, x3_pt).to(torch.float32)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y)
//  [python]
def test_input_tensor_constant(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = 2
        X3 = gen_input_tensor([dim, M], name="X3", dtype=dtype)
        Y = ops.where()(X1, X2, X3)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_input_tensor_constant")

        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = 2
            x3_pt = get_random_torch_tensor([batch, M], dtype)
            y_pt = torch.where(x1_pt, x2_pt, x3_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X3": x3_pt}
            model.run_with_tensors(inputs, [y])
            t.assert_close(y_pt, y)
//  [python]
def test_input_tensor_constant(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = 2
        X3 = gen_input_tensor([dim, M], name="X3", dtype=dtype)
        Y = ops.where()(X1, X2, X3)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_input_tensor_constant")

        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = 2
            x3_pt = get_random_torch_tensor([batch, M], dtype)
            y_pt = torch.where(x1_pt, x2_pt, x3_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X3": x3_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y)
//  [python]
def assert_models_equal(self, model_1, model_2, atol=1e-7, rtol=1e-7):
        assert model_1 is not model_2, 'Same model should not be compared.'
        for param1, param2 in zip(model_1.parameters(), model_2.parameters()):
            t.assert_close(param1, param2, atol=atol, rtol=rtol)
//  [python]
def assert_models_equal(self, model_1, model_2, atol=1e-7, rtol=1e-7):
        assert model_1 is not model_2, 'Same model should not be compared.'
        for param1, param2 in zip(model_1.parameters(), model_2.parameters()):
            torch.testing.assert_close(param1, param2, atol=atol, rtol=rtol)
//  [python]
def test_xauc(self) -> None:
        model_output = generate_model_output()
        self.xauc.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
        )
        metric = self.xauc.compute()
        actual_metric = metric[f"xauc-{DefaultTaskInfo.name}|lifetime_xauc"]
        expected_metric = model_output["expected_xauc"]

        t.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_xauc(self) -> None:
        model_output = generate_model_output()
        self.xauc.update(
            predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
            labels={DefaultTaskInfo.name: model_output["labels"][0]},
            weights={DefaultTaskInfo.name: model_output["weights"][0]},
        )
        metric = self.xauc.compute()
        actual_metric = metric[f"xauc-{DefaultTaskInfo.name}|lifetime_xauc"]
        expected_metric = model_output["expected_xauc"]

        torch.testing.assert_close(
            actual_metric,
            expected_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {expected_metric}",
        )
//  [python]
def test_where(self, dtype: str, M: int) -> None:
        dim = IntVar([2, 3, 128])
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = gen_input_tensor([dim, M], name="X2", dtype=dtype)
        X3 = gen_input_tensor([dim, M], name="X3", dtype=dtype)
        Y = ops.where()(X1, X2, X3)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", f"test_where_{self._test_id}")
        self._test_id += 1
        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = get_random_torch_tensor([batch, M], dtype)
            x3_pt = get_random_torch_tensor([batch, M], dtype)
            y_pt = torch.where(x1_pt, x2_pt, x3_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X2": x2_pt, "X3": x3_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(y_pt, y)
//  [python]
def test_where(self, dtype: str, M: int) -> None:
        dim = IntVar([2, 3, 128])
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = gen_input_tensor([dim, M], name="X2", dtype=dtype)
        X3 = gen_input_tensor([dim, M], name="X3", dtype=dtype)
        Y = ops.where()(X1, X2, X3)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", f"test_where_{self._test_id}")
        self._test_id += 1
        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = get_random_torch_tensor([batch, M], dtype)
            x3_pt = get_random_torch_tensor([batch, M], dtype)
            y_pt = torch.where(x1_pt, x2_pt, x3_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X2": x2_pt, "X3": x3_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y)
//  [python]
def test_other_tensor_constant(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = gen_input_tensor([dim, M], name="X2", dtype=dtype)
        X3 = 2
        Y = ops.where()(X1, X2, X3)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_other_tensor_constant")

        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = get_random_torch_tensor([batch, M], dtype)
            x3_pt = 2
            y_pt = torch.where(x1_pt, x2_pt, x3_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X2": x2_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(y_pt, y)
//  [python]
def test_other_tensor_constant(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype="bool")
        X2 = gen_input_tensor([dim, M], name="X2", dtype=dtype)
        X3 = 2
        Y = ops.where()(X1, X2, X3)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_other_tensor_constant")

        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype) < 0
            x2_pt = get_random_torch_tensor([batch, M], dtype)
            x3_pt = 2
            y_pt = torch.where(x1_pt, x2_pt, x3_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X2": x2_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y)
//  [python]
def assert_models_equal(self, model_1, model_2, threshold=None):
        if threshold is None:
            threshold = self.default_threshold

        assert model_1 is not model_2, 'Same model should not be compared.'
        for param1, param2 in zip(model_1.parameters(), model_2.parameters()):
            t.assert_close(param1, param2, **threshold)
//  [python]
def assert_models_equal(self, model_1, model_2, threshold=None):
        if threshold is None:
            threshold = self.default_threshold

        assert model_1 is not model_2, 'Same model should not be compared.'
        for param1, param2 in zip(model_1.parameters(), model_2.parameters()):
            torch.testing.assert_close(param1, param2, **threshold)
//  [python]
def test_integration_with_relational(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype=dtype)
        X2 = gen_input_tensor([dim, M], name="X2", dtype=dtype)
        X3 = ops.ge()(X1, X2)
        Y = ops.where()(X3, X1, X2)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_integration_with_relational")
        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype)
            x2_pt = get_random_torch_tensor([batch, M], dtype)
            x3_pt = torch.ge(x1_pt, x2_pt)
            y_pt = torch.where(x3_pt, x1_pt, x2_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X2": x2_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(y_pt, y)
//  [python]
def test_integration_with_relational(self) -> None:
        dim = IntVar([2, 3, 128])
        dtype = "float"
        M = 4
        X1 = gen_input_tensor([dim, M], name="X1", dtype=dtype)
        X2 = gen_input_tensor([dim, M], name="X2", dtype=dtype)
        X3 = ops.ge()(X1, X2)
        Y = ops.where()(X3, X1, X2)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        model = compile_model(Y, target, "./tmp", "test_integration_with_relational")
        for batch in dim._attrs["values"]:
            x1_pt = get_random_torch_tensor([batch, M], dtype)
            x2_pt = get_random_torch_tensor([batch, M], dtype)
            x3_pt = torch.ge(x1_pt, x2_pt)
            y_pt = torch.where(x3_pt, x1_pt, x2_pt)
            y = get_torch_empty_tensor([batch, M], dtype)
            inputs = {"X1": x1_pt, "X2": x2_pt}
            model.run_with_tensors(inputs, [y])
            torch.testing.assert_close(y_pt, y)
//  [python]
def test_log_survival_function():
    rng = numpy.random.default_rng(0xFEEEFEEE)

    batch_size = 32

    loc = rng.standard_normal(batch_size, dtype=numpy.float64)

    scale = rng.random(batch_size, dtype=numpy.float64) + 1.0

    x = numpy.linspace(-8.0, 8.0, batch_size, dtype=numpy.float64)

    t.assert_close(
        log_survival_function(
            torch.tensor(x, dtype=torch.float),
            Normal(
                torch.tensor(loc, dtype=torch.float),
                torch.tensor(scale, dtype=torch.float),
            ),
        ),
        torch.tensor(scipy.stats.norm(loc, scale).logsf(x), dtype=torch.float),
    )
//  [python]
def test_log_survival_function():
    rng = numpy.random.default_rng(0xFEEEFEEE)

    batch_size = 32

    loc = rng.standard_normal(batch_size, dtype=numpy.float64)

    scale = rng.random(batch_size, dtype=numpy.float64) + 1.0

    x = numpy.linspace(-8.0, 8.0, batch_size, dtype=numpy.float64)

    torch.testing.assert_close(
        log_survival_function(
            torch.tensor(x, dtype=torch.float),
            Normal(
                torch.tensor(loc, dtype=torch.float),
                torch.tensor(scale, dtype=torch.float),
            ),
        ),
        torch.tensor(scipy.stats.norm(loc, scale).logsf(x), dtype=torch.float),
    )
//  [python]
def vanilla_quantize_all_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, and quantize the module
    q_util = SelectiveQuantizer()
    q_util.quantize_module(module)

    print(module)  # You should expect to see QuantConv2d

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_allclose(y.size(), (1, 8, 32, 32))
//  [python]
def vanilla_quantize_all_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, and quantize the module
    q_util = SelectiveQuantizer()
    q_util.quantize_module(module)

    print(module)  # You should expect to see QuantConv2d

    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_shampoo_baseline_and_pt2(
        baseline_optim_factory: Callable[[Iterable[Parameter]], torch.optim.Optimizer],
        pt2_optim_factory: Callable[[Iterable[Parameter]], torch.optim.Optimizer],
        total_steps: int = 5,
    ) -> None:
        (
            baseline_params,
            baseline_loss,
        ) = DistributedShampooPytorchCompileTest._train_quadratic(
            baseline_optim_factory,
            total_steps=total_steps,
        )
        pt2_params, pt2_loss = DistributedShampooPytorchCompileTest._train_quadratic(
            pt2_optim_factory,
            total_steps=total_steps,
        )
        torch.testing.assert_allclose(pt2_loss, baseline_loss)
        torch.testing.assert_allclose(
            pt2_params,
            baseline_params,
        )
//  [python]
def _test_shampoo_baseline_and_pt2(
        baseline_optim_factory: Callable[[Iterable[Parameter]], torch.optim.Optimizer],
        pt2_optim_factory: Callable[[Iterable[Parameter]], torch.optim.Optimizer],
        total_steps: int = 5,
    ) -> None:
        (
            baseline_params,
            baseline_loss,
        ) = DistributedShampooPytorchCompileTest._train_quadratic(
            baseline_optim_factory,
            total_steps=total_steps,
        )
        pt2_params, pt2_loss = DistributedShampooPytorchCompileTest._train_quadratic(
            pt2_optim_factory,
            total_steps=total_steps,
        )
        torch.testing.assert_close(pt2_loss, baseline_loss)
        torch.testing.assert_close(
            pt2_params,
            baseline_params,
        )
//  [python]
def _test_and_verify(
        self,
        module,
        torch_output,
        dtype,
        has_d1=False,
        module_output_name="output_0",
    ):
        inputs = {
            "input_0": self.X_pt,
            "input_1": self.W_pt,
            "input_2": self.B_pt,
            "d0": self.D0_pt,
        }
        if has_d1:
            inputs["d1"] = self.D1_pt
        y = get_torch_empty_tensor(list(torch_output.shape), dtype)
        module.run_with_tensors(inputs, [y])
        if self.X_pt.nelement() == 0 or self.W_pt.nelement() == 0:
            pass
        else:
            t.assert_close(torch_output, y, atol=1e-1, rtol=1e-1)
//  [python]
def _test_and_verify(
        self,
        module,
        torch_output,
        dtype,
        has_d1=False,
        module_output_name="output_0",
    ):
        inputs = {
            "input_0": self.X_pt,
            "input_1": self.W_pt,
            "input_2": self.B_pt,
            "d0": self.D0_pt,
        }
        if has_d1:
            inputs["d1"] = self.D1_pt
        y = get_torch_empty_tensor(list(torch_output.shape), dtype)
        module.run_with_tensors(inputs, [y])
        if self.X_pt.nelement() == 0 or self.W_pt.nelement() == 0:
            pass
        else:
            torch.testing.assert_close(torch_output, y, atol=1e-1, rtol=1e-1)
//  [python]
def assert_close(expected, actual) -> None:
    if isinstance(expected, KeyedTensor):
        assert isinstance(actual, KeyedTensor)
        assert len(expected.keys()) == len(actual.keys())
        t.assert_close(expected.values(), actual.values())
        t.assert_close(expected.length_per_key(), actual.length_per_key())
    elif isinstance(expected, dict):
        assert sorted(expected.keys()) == sorted(actual.keys())
        for feature, jt_e in expected.items():
            jt_got = actual[feature]
            assert_close(jt_e.lengths(), jt_got.lengths())
            assert_close(jt_e.values(), jt_got.values())
            assert_close(jt_e.offsets(), jt_got.offsets())
    else:
        if isinstance(expected, torch.Tensor) and isinstance(actual, torch.Tensor):
            if actual.device != expected.device:
                actual = actual.to(expected.device)

        t.assert_close(expected, actual)
//  [python]
def assert_close(expected, actual) -> None:
    if isinstance(expected, KeyedTensor):
        assert isinstance(actual, KeyedTensor)
        assert len(expected.keys()) == len(actual.keys())
        torch.testing.assert_close(expected.values(), actual.values())
        torch.testing.assert_close(expected.length_per_key(), actual.length_per_key())
    elif isinstance(expected, dict):
        assert sorted(expected.keys()) == sorted(actual.keys())
        for feature, jt_e in expected.items():
            jt_got = actual[feature]
            assert_close(jt_e.lengths(), jt_got.lengths())
            assert_close(jt_e.values(), jt_got.values())
            assert_close(jt_e.offsets(), jt_got.offsets())
    else:
        if isinstance(expected, torch.Tensor) and isinstance(actual, torch.Tensor):
            if actual.device != expected.device:
                actual = actual.to(expected.device)

        torch.testing.assert_close(expected, actual)
//  [python]
def test_getitem_with_transforms(self, array_dataset):
        assert len(array_dataset[0]) == 4
        assert len(array_dataset[1]) == 4
        assert len(array_dataset[2]) == 4
        assert len(array_dataset[3]) == 4
        t.assert_close(array_dataset[0][0], torch.tensor([1, 0, -1, 2]))
        t.assert_close(array_dataset[0][1], torch.tensor(1))
        np.testing.assert_array_equal(array_dataset[0][2], np.array([2, 1, -5, 1]))
        assert array_dataset[0][3] == 1
        t.assert_close(array_dataset[1][0], torch.tensor([1, 0, -2, -1]))
        t.assert_close(array_dataset[1][1], torch.tensor(0))
        np.testing.assert_array_equal(array_dataset[1][2], np.array([1, 0, -2, -1]))
        assert array_dataset[1][3] == 0
        t.assert_close(array_dataset[2][0], torch.tensor([2, 5, 0, 3]))
        t.assert_close(array_dataset[2][1], torch.tensor(0))
        np.testing.assert_array_equal(array_dataset[2][2], np.array([2, 5, 0, 3]))
        assert array_dataset[2][3] == 1
        t.assert_close(array_dataset[3][0], torch.tensor([-7, 1, 2, 2]))
        t.assert_close(array_dataset[3][1], torch.tensor(1))
        np.testing.assert_array_equal(array_dataset[3][2], np.array([-7, 1, 2, 2]))
        assert array_dataset[3][3] == 1