jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_torch_slice(self):
        A = torch.randn((10, 5))
        tensors = {
            "a": A,
        }
        save_file_pt(tensors, "./slice.safetensors")

        # Now loading
        with safe_open("./slice.safetensors", framework="pt", device="cpu") as f:
            slice_ = f.get_slice("a")
            tensor = slice_[:]
            self.assertEqual(list(tensor.shape), [10, 5])
            torch.testing.assert_close(tensor, A)

            tensor = slice_[:2]
            self.assertEqual(list(tensor.shape), [2, 5])
            torch.testing.assert_close(tensor, A[:2])

            tensor = slice_[:, :2]
            self.assertEqual(list(tensor.shape), [10, 2])
            torch.testing.assert_close(tensor, A[:, :2])

            tensor = slice_[0, :2]
            self.assertEqual(list(tensor.shape), [2])
            torch.testing.assert_close(tensor, A[0, :2])

            tensor = slice_[2:, 0]
            self.assertEqual(list(tensor.shape), [8])
            torch.testing.assert_close(tensor, A[2:, 0])

            tensor = slice_[2:, 1]
            self.assertEqual(list(tensor.shape), [8])
            torch.testing.assert_close(tensor, A[2:, 1])

            tensor = slice_[2:, -1]
            self.assertEqual(list(tensor.shape), [8])
            torch.testing.assert_close(tensor, A[2:, -1])
//  [python]
def test_load_balancing_loss(self):
        r"""
        Let's make sure we can actually compute the loss and do a backward on it.
        """
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        config.num_local_experts = 8
        config.output_router_logits = True
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        model = MixtralForCausalLM(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask)
        self.assertEqual(result.router_logits[0].shape, (91, config.num_local_experts))
        self.assertTrue(torch.allclose(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2))

        # First, we make sure that adding padding tokens doesn't change the loss
        # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
        pad_length = 1000
        # Add padding tokens (assume that pad_token_id=1) to input_ids
        padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
        padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left
        padded_attention_mask = padded_input_ids.ne(1).to(torch_device)

        padded_result = model(padded_input_ids, attention_mask=padded_attention_mask)
        self.assertTrue(torch.allclose(result.aux_loss.cpu(), padded_result.aux_loss.cpu(), rtol=1e-4, atol=1e-4))

        # We make sure that the loss of includding padding tokens != the loss without padding tokens
        # if attention_mask=None --> we don't exclude padding tokens
        include_padding_result = model(padded_input_ids, attention_mask=None)

        # This is to mimic torch.testing.assert_not_close
        self.assertNotAlmostEqual(include_padding_result.aux_loss.item(), result.aux_loss.item())
//  [python]
def test_load_balancing_loss(self):
        r"""
        Let's make sure we can actually compute the loss and do a backward on it.
        """
        config, input_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.num_labels = 3
        config.num_local_experts = 8
        config.output_router_logits = True
        input_ids = input_dict["input_ids"]
        attention_mask = input_ids.ne(1).to(torch_device)
        model = MixtralForCausalLM(config)
        model.to(torch_device)
        model.eval()
        result = model(input_ids, attention_mask=attention_mask)
        self.assertEqual(result.router_logits[0].shape, (91, config.num_local_experts))
        torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(2, dtype=torch.float32), rtol=1e-2, atol=1e-2)

        # First, we make sure that adding padding tokens doesn't change the loss
        # loss(input_ids, attention_mask=None) == loss(input_ids + padding, attention_mask=attention_mask_with_padding)
        pad_length = 1000
        # Add padding tokens (assume that pad_token_id=1) to input_ids
        padding_block = torch.ones(input_ids.shape[0], pad_length, dtype=torch.int32).to(torch_device)
        padded_input_ids = torch.cat((padding_block, input_ids), dim=1)  # this is to simulate padding to the left
        padded_attention_mask = padded_input_ids.ne(1).to(torch_device)

        padded_result = model(padded_input_ids, attention_mask=padded_attention_mask)
        torch.testing.assert_close(result.aux_loss.cpu(), padded_result.aux_loss.cpu(), rtol=1e-4, atol=1e-4)

        # We make sure that the loss of includding padding tokens != the loss without padding tokens
        # if attention_mask=None --> we don't exclude padding tokens
        include_padding_result = model(padded_input_ids, attention_mask=None)

        # This is to mimic torch.testing.assert_not_close
        self.assertNotAlmostEqual(include_padding_result.aux_loss.item(), result.aux_loss.item())
//  [python]
def test_all_augment(self, inputs):
        aug = RandAug(num_ops=15, prob=1, sample_with_replacement=False)
        actual_img = aug(inputs)
        t.assert_allclose(
            transforms.ToTensor()(actual_img),
            torch.Tensor(
                [
                    [
                        [0.7490, 0.0118, 0.3373, 0.8745],
                        [0.5333, 0.0667, 0.4667, 0.0275],
                        [0.5098, 0.0510, 0.5608, 0.6627],
                        [0.8039, 0.0000, 0.4863, 0.4863],
                    ],
                    [
                        [0.7843, 0.0118, 0.3294, 0.8784],
                        [0.5490, 0.0627, 0.4667, 0.0235],
                        [0.5294, 0.0471, 0.5647, 0.6941],
                        [0.8392, 0.0000, 0.4549, 0.4549],
                    ],
                    [
                        [0.2000, 0.0118, 0.3294, 0.7882],
                        [0.6667, 0.0667, 0.4667, 0.0980],
                        [0.4784, 0.0471, 0.5608, 0.1098],
                        [0.1020, 0.0000, 0.4078, 0.4078],
                    ],
                ]
            ),
            atol=0.0001,
            rtol=0.0,
        )
//  [python]
def test_all_augment(self, inputs):
        aug = RandAug(num_ops=15, prob=1, sample_with_replacement=False)
        actual_img = aug(inputs)
        torch.testing.assert_close(
            transforms.ToTensor()(actual_img),
            torch.Tensor(
                [
                    [
                        [0.7490, 0.0118, 0.3373, 0.8745],
                        [0.5333, 0.0667, 0.4667, 0.0275],
                        [0.5098, 0.0510, 0.5608, 0.6627],
                        [0.8039, 0.0000, 0.4863, 0.4863],
                    ],
                    [
                        [0.7843, 0.0118, 0.3294, 0.8784],
                        [0.5490, 0.0627, 0.4667, 0.0235],
                        [0.5294, 0.0471, 0.5647, 0.6941],
                        [0.8392, 0.0000, 0.4549, 0.4549],
                    ],
                    [
                        [0.2000, 0.0118, 0.3294, 0.7882],
                        [0.6667, 0.0667, 0.4667, 0.0980],
                        [0.4784, 0.0471, 0.5608, 0.1098],
                        [0.1020, 0.0000, 0.4078, 0.4078],
                    ],
                ]
            ),
            atol=0.0001,
            rtol=0.0,
        )
//  [python]
def test_quat_unique(self):
        """Test quat_unique method."""
        # Define test cases
        quats = math_utils.random_orientation(num=1024, device="cpu")

        # Test positive real quaternion
        pos_real_quats = math_utils.quat_unique(quats)

        # Test that the real part is positive
        self.assertTrue(torch.all(pos_real_quats[:, 0] > 0).item())

        non_pos_indices = quats[:, 0] < 0
        # Check imaginary part have sign flipped if real part is negative
        self.assertEqual(pos_real_quats[non_pos_indices], -quats[non_pos_indices])
        self.assertEqual(pos_real_quats[~non_pos_indices], quats[~non_pos_indices])
//  [python]
def test_quat_unique(self):
        """Test quat_unique method."""
        # Define test cases
        quats = math_utils.random_orientation(num=1024, device="cpu")

        # Test positive real quaternion
        pos_real_quats = math_utils.quat_unique(quats)

        # Test that the real part is positive
        self.assertTrue(torch.all(pos_real_quats[:, 0] > 0).item())

        non_pos_indices = quats[:, 0] < 0
        # Check imaginary part have sign flipped if real part is negative
        torch.testing.assert_close(pos_real_quats[non_pos_indices], -quats[non_pos_indices])
        torch.testing.assert_close(pos_real_quats[~non_pos_indices], quats[~non_pos_indices])
//  [python]
def test_axis_angle_from_quat_approximation(self):
        """Test the Taylor approximation from axis_angle_from_quat method.

        This test checks for unstable conversions where theta is very small.
        """
        # Generate a small rotation quaternion
        # Small angle
        theta = torch.Tensor([0.0000001])
        # Arbitrary normalized axis of rotation in rads, (x,y,z)
        axis = [-0.302286, 0.205494, -0.930803]
        # Generate quaternion
        qw = torch.cos(theta / 2)
        quat_vect = [qw] + [d * torch.sin(theta / 2) for d in axis]
        quaternion = torch.tensor(quat_vect, dtype=torch.float32)

        # Convert quaternion to axis-angle
        axis_angle_computed = math_utils.axis_angle_from_quat(quaternion)

        # Expected axis-angle representation
        axis_angle_expected = torch.tensor([theta * d for d in axis], dtype=torch.float32)

        # Assert that the computed values are close to the expected values
        torch.testing.assert_allclose(axis_angle_computed, axis_angle_expected)
//  [python]
def test_axis_angle_from_quat_approximation(self):
        """Test the Taylor approximation from axis_angle_from_quat method.

        This test checks for unstable conversions where theta is very small.
        """
        # Generate a small rotation quaternion
        # Small angle
        theta = torch.Tensor([0.0000001])
        # Arbitrary normalized axis of rotation in rads, (x,y,z)
        axis = [-0.302286, 0.205494, -0.930803]
        # Generate quaternion
        qw = torch.cos(theta / 2)
        quat_vect = [qw] + [d * torch.sin(theta / 2) for d in axis]
        quaternion = torch.tensor(quat_vect, dtype=torch.float32)

        # Convert quaternion to axis-angle
        axis_angle_computed = math_utils.axis_angle_from_quat(quaternion)

        # Expected axis-angle representation
        axis_angle_expected = torch.tensor([theta * d for d in axis], dtype=torch.float32)

        # Assert that the computed values are close to the expected values
        torch.testing.assert_close(axis_angle_computed, axis_angle_expected)
//  [python]
def test_axis_angle_from_quat(self):
        """Test axis_angle_from_quat method."""
        # Quaternions of the form (2,4) and (2,2,4)
        quats = [
            torch.Tensor([[1.0, 0.0, 0.0, 0.0], [0.8418536, 0.142006, 0.0, 0.5206887]]),
            torch.Tensor([
                [[1.0, 0.0, 0.0, 0.0], [0.8418536, 0.142006, 0.0, 0.5206887]],
                [[1.0, 0.0, 0.0, 0.0], [0.9850375, 0.0995007, 0.0995007, 0.0995007]],
            ]),
        ]

        # Angles of the form (2,3) and (2,2,3)
        angles = [
            torch.Tensor([[0.0, 0.0, 0.0], [0.3, 0.0, 1.1]]),
            torch.Tensor([[[0.0, 0.0, 0.0], [0.3, 0.0, 1.1]], [[0.0, 0.0, 0.0], [0.2, 0.2, 0.2]]]),
        ]

        for quat, angle in zip(quats, angles):
            with self.subTest(quat=quat, angle=angle):
                torch.testing.assert_allclose(math_utils.axis_angle_from_quat(quat), angle)
//  [python]
def test_axis_angle_from_quat(self):
        """Test axis_angle_from_quat method."""
        # Quaternions of the form (2,4) and (2,2,4)
        quats = [
            torch.Tensor([[1.0, 0.0, 0.0, 0.0], [0.8418536, 0.142006, 0.0, 0.5206887]]),
            torch.Tensor([
                [[1.0, 0.0, 0.0, 0.0], [0.8418536, 0.142006, 0.0, 0.5206887]],
                [[1.0, 0.0, 0.0, 0.0], [0.9850375, 0.0995007, 0.0995007, 0.0995007]],
            ]),
        ]

        # Angles of the form (2,3) and (2,2,3)
        angles = [
            torch.Tensor([[0.0, 0.0, 0.0], [0.3, 0.0, 1.1]]),
            torch.Tensor([[[0.0, 0.0, 0.0], [0.3, 0.0, 1.1]], [[0.0, 0.0, 0.0], [0.2, 0.2, 0.2]]]),
        ]

        for quat, angle in zip(quats, angles):
            with self.subTest(quat=quat, angle=angle):
                torch.testing.assert_close(math_utils.axis_angle_from_quat(quat), angle)
//  [python]
def test_quat_error_magnitude(self):
        """Test quat_error_magnitude method."""
        # Define test cases
        # Each tuple contains: q1, q2, expected error
        test_cases = [
            # No rotation
            (torch.Tensor([1, 0, 0, 0]), torch.Tensor([1, 0, 0, 0]), torch.Tensor([0.0])),
            # PI/2 rotation
            (torch.Tensor([1.0, 0, 0.0, 0]), torch.Tensor([0.7071068, 0.7071068, 0, 0]), torch.Tensor([PI / 2])),
            # PI rotation
            (torch.Tensor([1.0, 0, 0.0, 0]), torch.Tensor([0.0, 0.0, 1.0, 0]), torch.Tensor([PI])),
        ]
        # Test higher dimension (batched) inputs
        q1_list = torch.stack([t[0] for t in test_cases], dim=0)
        q2_list = torch.stack([t[1] for t in test_cases], dim=0)
        expected_diff_list = torch.stack([t[2] for t in test_cases], dim=0).flatten()
        test_cases += [(q1_list, q2_list, expected_diff_list)]

        # Iterate over test cases
        for q1, q2, expected_diff in test_cases:
            with self.subTest(q1=q1, q2=q2):
                # Compute the error
                q12_diff = math_utils.quat_error_magnitude(q1, q2)

                # Check that the error is close to the expected value
                if len(q1.shape) > 1:
                    self.assertTrue(torch.allclose(q12_diff, expected_diff))
                else:
                    self.assertAlmostEqual(q12_diff.item(), expected_diff.item(), places=5)
//  [python]
def test_quat_error_magnitude(self):
        """Test quat_error_magnitude method."""
        # Define test cases
        # Each tuple contains: q1, q2, expected error
        test_cases = [
            # No rotation
            (torch.Tensor([1, 0, 0, 0]), torch.Tensor([1, 0, 0, 0]), torch.Tensor([0.0])),
            # PI/2 rotation
            (torch.Tensor([1.0, 0, 0.0, 0]), torch.Tensor([0.7071068, 0.7071068, 0, 0]), torch.Tensor([PI / 2])),
            # PI rotation
            (torch.Tensor([1.0, 0, 0.0, 0]), torch.Tensor([0.0, 0.0, 1.0, 0]), torch.Tensor([PI])),
        ]
        # Test higher dimension (batched) inputs
        q1_list = torch.stack([t[0] for t in test_cases], dim=0)
        q2_list = torch.stack([t[1] for t in test_cases], dim=0)
        expected_diff_list = torch.stack([t[2] for t in test_cases], dim=0).flatten()
        test_cases += [(q1_list, q2_list, expected_diff_list)]

        # Iterate over test cases
        for q1, q2, expected_diff in test_cases:
            with self.subTest(q1=q1, q2=q2):
                # Compute the error
                q12_diff = math_utils.quat_error_magnitude(q1, q2)

                # Check that the error is close to the expected value
                if len(q1.shape) > 1:
                    torch.testing.assert_close(q12_diff, expected_diff)
                else:
                    self.assertAlmostEqual(q12_diff.item(), expected_diff.item(), places=5)
//  [python]
def _make_model_unit_test(
        self,
        model,
        *args,
        input_shape=None,
        enable_allclose=False,
        **kwargs,
    ):
        """
        Test that the model can be traced correctly and is producing correct
        result.
        """
        if input_shape is None:
            input_shape = [1, 3, 224, 224]
        input = torch.randn(input_shape)
        traced = acc_tracer.trace(model, [input])
        if enable_allclose:
            self.assertEqual(model(input), traced(input))
        else:
            self.assertTrue(torch.equal(model(input), traced(input)))
        traced_again = acc_tracer.trace(traced, [input])
        if enable_allclose:
            self.assertEqual(model(input), traced_again(input))
        else:
            self.assertTrue(torch.equal(model(input), traced_again(input)))
//  [python]
def _make_model_unit_test(
        self,
        model,
        *args,
        input_shape=None,
        enable_allclose=False,
        **kwargs,
    ):
        """
        Test that the model can be traced correctly and is producing correct
        result.
        """
        if input_shape is None:
            input_shape = [1, 3, 224, 224]
        input = torch.randn(input_shape)
        traced = acc_tracer.trace(model, [input])
        if enable_allclose:
            torch.testing.assert_close(model(input), traced(input))
        else:
            self.assertTrue(torch.equal(model(input), traced(input)))
        traced_again = acc_tracer.trace(traced, [input])
        if enable_allclose:
            torch.testing.assert_close(model(input), traced_again(input))
        else:
            self.assertTrue(torch.equal(model(input), traced_again(input)))
//  [python]
def test_video_recording_move_to_memory(video_recording):
    memory_recording = video_recording.move_to_memory()

    video, audio = video_recording.load_video()
    video1, audio1 = memory_recording.load_video()

    t.assert_close(video, video1)
    t.assert_close(audio, audio1)
//  [python]
def test_video_recording_move_to_memory(video_recording):
    memory_recording = video_recording.move_to_memory()

    video, audio = video_recording.load_video()
    video1, audio1 = memory_recording.load_video()

    torch.testing.assert_close(video, video1)
    torch.testing.assert_close(audio, audio1)
//  [python]
def test_video_recording_load_video_duration(video_recording):
    video, _ = video_recording.load_video(duration=2.0, with_audio=False)

    expected_dims = (25 * 2, COLOR, HEIGHT, WIDTH)
    assert video.shape == expected_dims

    full_video, _ = video_recording.load_video(with_audio=False)

    t.assert_allclose(video, full_video[:50])
//  [python]
def test_video_recording_load_video_duration(video_recording):
    video, _ = video_recording.load_video(duration=2.0, with_audio=False)

    expected_dims = (25 * 2, COLOR, HEIGHT, WIDTH)
    assert video.shape == expected_dims

    full_video, _ = video_recording.load_video(with_audio=False)

    torch.testing.assert_close(video, full_video[:50])
//  [python]
def test_quat_error_mag_with_quat_unique(self):
        """Test quat_error_magnitude method with positive real quaternions."""

        quats_1 = math_utils.random_orientation(num=1024, device="cpu")
        quats_2 = math_utils.random_orientation(num=1024, device="cpu")
        # Make quats positive real
        quats_1_pos_real = math_utils.quat_unique(quats_1)
        quats_2_pos_real = math_utils.quat_unique(quats_2)

        # Compute the error
        error_1 = math_utils.quat_error_magnitude(quats_1, quats_2)
        error_2 = math_utils.quat_error_magnitude(quats_1_pos_real, quats_2_pos_real)
        error_3 = math_utils.quat_error_magnitude(quats_1, quats_2_pos_real)
        error_4 = math_utils.quat_error_magnitude(quats_1_pos_real, quats_2)

        # Check that the error is close to the expected value
        t.assert_close(error_1, error_2)
        t.assert_close(error_2, error_3)
        t.assert_close(error_3, error_4)
        t.assert_close(error_4, error_1)
//  [python]
def test_quat_error_mag_with_quat_unique(self):
        """Test quat_error_magnitude method with positive real quaternions."""

        quats_1 = math_utils.random_orientation(num=1024, device="cpu")
        quats_2 = math_utils.random_orientation(num=1024, device="cpu")
        # Make quats positive real
        quats_1_pos_real = math_utils.quat_unique(quats_1)
        quats_2_pos_real = math_utils.quat_unique(quats_2)

        # Compute the error
        error_1 = math_utils.quat_error_magnitude(quats_1, quats_2)
        error_2 = math_utils.quat_error_magnitude(quats_1_pos_real, quats_2_pos_real)
        error_3 = math_utils.quat_error_magnitude(quats_1, quats_2_pos_real)
        error_4 = math_utils.quat_error_magnitude(quats_1_pos_real, quats_2)

        # Check that the error is close to the expected value
        torch.testing.assert_close(error_1, error_2)
        torch.testing.assert_close(error_2, error_3)
        torch.testing.assert_close(error_3, error_4)
        torch.testing.assert_close(error_4, error_1)
//  [python]
def test_quat_mul_with_quat_unique(self):
        """Test quat_mul method with different quaternions.

        This test checks that the quaternion multiplication is consistent when using positive real quaternions
        and regular quaternions. It makes sure that the result is the same regardless of the input quaternion sign
        (i.e. q and -q are same quaternion in the context of rotations).
        """

        quats_1 = math_utils.random_orientation(num=1024, device="cpu")
        quats_2 = math_utils.random_orientation(num=1024, device="cpu")
        # Make quats positive real
        quats_1_pos_real = math_utils.quat_unique(quats_1)
        quats_2_pos_real = math_utils.quat_unique(quats_2)

        # Option 1: Direct computation on quaternions
        quat_result_1 = math_utils.quat_mul(quats_1, math_utils.quat_conjugate(quats_2))
        quat_result_1 = math_utils.quat_unique(quat_result_1)

        # Option 2: Computation on positive real quaternions
        quat_result_2 = math_utils.quat_mul(quats_1_pos_real, math_utils.quat_conjugate(quats_2_pos_real))
        quat_result_2 = math_utils.quat_unique(quat_result_2)

        # Option 3: Mixed computation
        quat_result_3 = math_utils.quat_mul(quats_1, math_utils.quat_conjugate(quats_2_pos_real))
        quat_result_3 = math_utils.quat_unique(quat_result_3)

        # Check that the result is close to the expected value
        self.assertTrue(torch.allclose(quat_result_1, quat_result_2))
        self.assertTrue(torch.allclose(quat_result_2, quat_result_3))
        self.assertTrue(torch.allclose(quat_result_3, quat_result_1))
//  [python]
def test_quat_mul_with_quat_unique(self):
        """Test quat_mul method with different quaternions.

        This test checks that the quaternion multiplication is consistent when using positive real quaternions
        and regular quaternions. It makes sure that the result is the same regardless of the input quaternion sign
        (i.e. q and -q are same quaternion in the context of rotations).
        """

        quats_1 = math_utils.random_orientation(num=1024, device="cpu")
        quats_2 = math_utils.random_orientation(num=1024, device="cpu")
        # Make quats positive real
        quats_1_pos_real = math_utils.quat_unique(quats_1)
        quats_2_pos_real = math_utils.quat_unique(quats_2)

        # Option 1: Direct computation on quaternions
        quat_result_1 = math_utils.quat_mul(quats_1, math_utils.quat_conjugate(quats_2))
        quat_result_1 = math_utils.quat_unique(quat_result_1)

        # Option 2: Computation on positive real quaternions
        quat_result_2 = math_utils.quat_mul(quats_1_pos_real, math_utils.quat_conjugate(quats_2_pos_real))
        quat_result_2 = math_utils.quat_unique(quat_result_2)

        # Option 3: Mixed computation
        quat_result_3 = math_utils.quat_mul(quats_1, math_utils.quat_conjugate(quats_2_pos_real))
        quat_result_3 = math_utils.quat_unique(quat_result_3)

        # Check that the result is close to the expected value
        torch.testing.assert_close(quat_result_1, quat_result_2)
        torch.testing.assert_close(quat_result_2, quat_result_3)
        torch.testing.assert_close(quat_result_3, quat_result_1)
//  [python]
def test_video_recording_load_video_over_duration(video_recording):
    video, _ = video_recording.load_video(duration=10.0, with_audio=False)

    expected_dims = (132, COLOR, HEIGHT, WIDTH)
    assert video.shape == expected_dims

    full_video, _ = video_recording.load_video(with_audio=False)

    t.assert_close(video, full_video)
//  [python]
def test_video_recording_load_video_over_duration(video_recording):
    video, _ = video_recording.load_video(duration=10.0, with_audio=False)

    expected_dims = (132, COLOR, HEIGHT, WIDTH)
    assert video.shape == expected_dims

    full_video, _ = video_recording.load_video(with_audio=False)

    torch.testing.assert_close(video, full_video)
//  [python]
def test_hardswish(self):
        class TestModule(nn.Module):
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                y = nn.functional.hardswish(x)
                return y

        m = TestModule()
        x = torch.randn(3, 4, 5)
        traced = acc_tracer.trace(m, [x])
        ph_x = hardsigmoid_y = res_y = None
        for node in traced.graph.nodes:
            if node.op == "placeholder":
                ph_x = node
            elif node.op == "call_function" and node.target == acc_ops.hardsigmoid:
                hardsigmoid_y = node
                self.assertEqual(node.kwargs["input"], ph_x)
            elif node.op == "call_function" and node.target == acc_ops.mul:
                res_y = node
                self.assertEqual(node.kwargs["input"], hardsigmoid_y)
                self.assertEqual(node.kwargs["other"], ph_x)
            elif node.op == "output":
                self.assertEqual(node.args[0], res_y)
            else:
                self.fail(f"Unexpected node: {node.format_node()}")

        ref = m(x)
        res = traced(x)
        t.assert_allclose(ref, res)
//  [python]
def test_hardswish(self):
        class TestModule(nn.Module):
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                y = nn.functional.hardswish(x)
                return y

        m = TestModule()
        x = torch.randn(3, 4, 5)
        traced = acc_tracer.trace(m, [x])
        ph_x = hardsigmoid_y = res_y = None
        for node in traced.graph.nodes:
            if node.op == "placeholder":
                ph_x = node
            elif node.op == "call_function" and node.target == acc_ops.hardsigmoid:
                hardsigmoid_y = node
                self.assertEqual(node.kwargs["input"], ph_x)
            elif node.op == "call_function" and node.target == acc_ops.mul:
                res_y = node
                self.assertEqual(node.kwargs["input"], hardsigmoid_y)
                self.assertEqual(node.kwargs["other"], ph_x)
            elif node.op == "output":
                self.assertEqual(node.args[0], res_y)
            else:
                self.fail(f"Unexpected node: {node.format_node()}")

        ref = m(x)
        res = traced(x)
        torch.testing.assert_close(ref, res)
//  [python]
def test_log1p(self):
        class TestModule(torch.nn.Module):
            def forward(self, input: torch.Tensor) -> torch.Tensor:
                return torch.log1p(input)

        m = TestModule().eval()
        input = torch.tensor([[1.2, 0.3, -0.4]])
        traced = acc_tracer.trace(m, [input])

        ph_in = add = log = None
        for node in traced.graph.nodes:
            if node.op == "placeholder":
                self.assertTrue(str(node.target) == "input")
                ph_in = node
            elif node.op == "call_function":
                if node.target == acc_ops.add:
                    self.assertEqual(node.kwargs["input"], ph_in)
                    self.assertEqual(node.kwargs["other"], 1)
                    add = node
                else:
                    self.assertEqual(node.target, acc_ops.log)
                    self.assertEqual(node.kwargs["input"], add)
                    log = node
            elif node.op == "output":
                self.assertEqual(log, node.args[0])
            else:
                self.fail(f"Unexpected node: {node.format_node()}")

        torch.testing.assert_allclose(m(input), traced(input))
//  [python]
def test_log1p(self):
        class TestModule(torch.nn.Module):
            def forward(self, input: torch.Tensor) -> torch.Tensor:
                return torch.log1p(input)

        m = TestModule().eval()
        input = torch.tensor([[1.2, 0.3, -0.4]])
        traced = acc_tracer.trace(m, [input])

        ph_in = add = log = None
        for node in traced.graph.nodes:
            if node.op == "placeholder":
                self.assertTrue(str(node.target) == "input")
                ph_in = node
            elif node.op == "call_function":
                if node.target == acc_ops.add:
                    self.assertEqual(node.kwargs["input"], ph_in)
                    self.assertEqual(node.kwargs["other"], 1)
                    add = node
                else:
                    self.assertEqual(node.target, acc_ops.log)
                    self.assertEqual(node.kwargs["input"], add)
                    log = node
            elif node.op == "output":
                self.assertEqual(log, node.args[0])
            else:
                self.fail(f"Unexpected node: {node.format_node()}")

        torch.testing.assert_close(m(input), traced(input))
//  [python]
def test_export_for_inference_torchscript(model_cls, sample_input):
    model = model_cls()

    model.eval()

    orig_out = model(sample_input)
    save_format = 'torchscript'
    with tempfile.TemporaryDirectory() as tempdir:
        save_path = os.path.join(tempdir, f'model.pt')
        inference.export_for_inference(
            model=model,
            save_format=save_format,
            save_path=save_path,
        )
        loaded_model = torch.jit.load(save_path)
        loaded_model.eval()
        loaded_model_out = loaded_model(sample_input)

        torch.testing.assert_allclose(
            orig_out,
            loaded_model_out,
            msg=f'output mismatch with {save_format}',
        )
//  [python]
def test_export_for_inference_torchscript(model_cls, sample_input):
    model = model_cls()

    model.eval()

    orig_out = model(sample_input)
    save_format = 'torchscript'
    with tempfile.TemporaryDirectory() as tempdir:
        save_path = os.path.join(tempdir, f'model.pt')
        inference.export_for_inference(
            model=model,
            save_format=save_format,
            save_path=save_path,
        )
        loaded_model = torch.jit.load(save_path)
        loaded_model.eval()
        loaded_model_out = loaded_model(sample_input)

        torch.testing.assert_close(
            orig_out,
            loaded_model_out,
            msg=f'output mismatch with {save_format}',
        )
//  [python]
def test_saved_net_is_same(
            self, net, data, mock_hf_api, hf_hub_storer_cls, storage, tmp_path
    ):
        # Check that the pickled net has the same params after loading, both for
        # in-memory and on disk
        from skorch.callbacks import TrainEndCheckpoint

        if storage == 'memory':
            local_storage = None
        elif storage == 'str':
            local_storage = str(tmp_path / 'my-net.pkl')
        else:
            local_storage = tmp_path / 'my-net.pkl'

        storer = hf_hub_storer_cls(
            mock_hf_api,
            path_in_repo='my-model',
            repo_id='my-user/my-repo',
            local_storage=local_storage,
        )
        checkpoint = TrainEndCheckpoint(
            f_pickle=storer,
            f_params=None,
            f_optimizer=None,
            f_criterion=None,
            f_history=None,
        )
        net.set_params(callbacks=[checkpoint])
        net.fit(*data)
        net_loaded = pickle.loads(mock_hf_api.saved.read())

        assert len(net_loaded.module_.state_dict()) == len(net.module_.state_dict())
        for key, original in net_loaded.module_.state_dict().items():
            original = net.module_.state_dict()[key]
            loaded = net_loaded.module_.state_dict()[key]
            torch.allclose(loaded, original)
//  [python]
def test_saved_net_is_same(
            self, net, data, mock_hf_api, hf_hub_storer_cls, storage, tmp_path
    ):
        # Check that the pickled net has the same params after loading, both for
        # in-memory and on disk
        from skorch.callbacks import TrainEndCheckpoint

        if storage == 'memory':
            local_storage = None
        elif storage == 'str':
            local_storage = str(tmp_path / 'my-net.pkl')
        else:
            local_storage = tmp_path / 'my-net.pkl'

        storer = hf_hub_storer_cls(
            mock_hf_api,
            path_in_repo='my-model',
            repo_id='my-user/my-repo',
            local_storage=local_storage,
        )
        checkpoint = TrainEndCheckpoint(
            f_pickle=storer,
            f_params=None,
            f_optimizer=None,
            f_criterion=None,
            f_history=None,
        )
        net.set_params(callbacks=[checkpoint])
        net.fit(*data)
        net_loaded = pickle.loads(mock_hf_api.saved.read())

        assert len(net_loaded.module_.state_dict()) == len(net.module_.state_dict())
        for key, original in net_loaded.module_.state_dict().items():
            original = net.module_.state_dict()[key]
            loaded = net_loaded.module_.state_dict()[key]
            torch.testing.assert_close(loaded, original)
//  [python]
def test_saved_params_is_same(
            self, net, data, mock_hf_api, hf_hub_storer_cls, storage, tmp_path
    ):
        # check that the module parameters are the same after loading, both for
        # in-memory and on disk
        from skorch.callbacks import TrainEndCheckpoint

        if storage == 'memory':
            local_storage = None
        elif storage == 'str':
            local_storage = str(tmp_path / 'my-weights.pt')
        else:
            local_storage = tmp_path / 'my-weights.pt'

        storer = hf_hub_storer_cls(
            mock_hf_api,
            path_in_repo='my-model',
            repo_id='my-user/my-repo',
            local_storage=local_storage,
        )
        checkpoint = TrainEndCheckpoint(
            f_params=storer,
            f_optimizer=None,
            f_criterion=None,
            f_history=None,
        )
        net.set_params(callbacks=[checkpoint])
        net.fit(*data)
        state_dict_before = net.module_.state_dict()
        state_dict_after = torch.load(mock_hf_api.saved)

        assert len(state_dict_before) == len(state_dict_after)
        for key, original in state_dict_before.items():
            loaded = state_dict_after[key]
            torch.testing.assert_allclose(loaded, original)
//  [python]
def test_saved_params_is_same(
            self, net, data, mock_hf_api, hf_hub_storer_cls, storage, tmp_path
    ):
        # check that the module parameters are the same after loading, both for
        # in-memory and on disk
        from skorch.callbacks import TrainEndCheckpoint

        if storage == 'memory':
            local_storage = None
        elif storage == 'str':
            local_storage = str(tmp_path / 'my-weights.pt')
        else:
            local_storage = tmp_path / 'my-weights.pt'

        storer = hf_hub_storer_cls(
            mock_hf_api,
            path_in_repo='my-model',
            repo_id='my-user/my-repo',
            local_storage=local_storage,
        )
        checkpoint = TrainEndCheckpoint(
            f_params=storer,
            f_optimizer=None,
            f_criterion=None,
            f_history=None,
        )
        net.set_params(callbacks=[checkpoint])
        net.fit(*data)
        state_dict_before = net.module_.state_dict()
        state_dict_after = torch.load(mock_hf_api.saved)

        assert len(state_dict_before) == len(state_dict_after)
        for key, original in state_dict_before.items():
            loaded = state_dict_after[key]
            torch.testing.assert_close(loaded, original)
//  [python]
def test_export_for_inference_torchscript_ddp(model_cls, sample_input, request: pytest.FixtureRequest):
    model = model_cls()
    optimizer = torch.optim.SGD(model.parameters(), 0.1)
    device = None
    for item in request.session.items:
        device = DeviceCPU() if item.get_closest_marker('gpu') is None else DeviceGPU()
        break
    assert device != None

    state = State(
        model=model,
        rank_zero_seed=0,
        device=device,
        run_name='run_name',
        optimizers=optimizer,
        max_duration='1ep',
        dataloader=DataLoader(RandomImageDataset(shape=(3, 224, 224))),
        dataloader_label='train',
        precision='fp32',
    )

    state.model = prepare_ddp_module(state.model, find_unused_parameters=True)
    state.model.eval()
    orig_out = state.model(sample_input)

    save_format = 'torchscript'

    # Only one rank needs to save/load model
    if dist.get_local_rank() == 0:
        with tempfile.TemporaryDirectory() as tempdir:
            save_path = os.path.join(str(tempdir), f'model.pt')
            assert isinstance(state.model.module, nn.Module)
            inference.export_for_inference(
                model=state.model.module,
                save_format=save_format,
                save_path=save_path,
            )

            loaded_model = torch.jit.load(save_path)
            loaded_model.eval()
            loaded_model_out = loaded_model(sample_input)

            torch.testing.assert_allclose(orig_out, loaded_model_out)
//  [python]
def test_export_for_inference_torchscript_ddp(model_cls, sample_input, request: pytest.FixtureRequest):
    model = model_cls()
    optimizer = torch.optim.SGD(model.parameters(), 0.1)
    device = None
    for item in request.session.items:
        device = DeviceCPU() if item.get_closest_marker('gpu') is None else DeviceGPU()
        break
    assert device != None

    state = State(
        model=model,
        rank_zero_seed=0,
        device=device,
        run_name='run_name',
        optimizers=optimizer,
        max_duration='1ep',
        dataloader=DataLoader(RandomImageDataset(shape=(3, 224, 224))),
        dataloader_label='train',
        precision='fp32',
    )

    state.model = prepare_ddp_module(state.model, find_unused_parameters=True)
    state.model.eval()
    orig_out = state.model(sample_input)

    save_format = 'torchscript'

    # Only one rank needs to save/load model
    if dist.get_local_rank() == 0:
        with tempfile.TemporaryDirectory() as tempdir:
            save_path = os.path.join(str(tempdir), f'model.pt')
            assert isinstance(state.model.module, nn.Module)
            inference.export_for_inference(
                model=state.model.module,
                save_format=save_format,
                save_path=save_path,
            )

            loaded_model = torch.jit.load(save_path)
            loaded_model.eval()
            loaded_model_out = loaded_model(sample_input)

            torch.testing.assert_close(orig_out, loaded_model_out)
//  [python]
def test_softmax(BLOCK, WIDTH, DTYPE):
    # set seed
    torch.random.manual_seed(0)
    Z, H, M, N = 2, 4, WIDTH, WIDTH
    scale = 0.4

    # create inputs
    layout = torch.randint(2, (H, M // BLOCK, N // BLOCK))
    x = torch.randn((Z, H, M, N), dtype=DTYPE, requires_grad=True, device="cuda")

    # triton result
    op = blocksparse_softmax(layout, BLOCK, device=torch.device("cuda"))
    tx = block_sparsify_tensor(x, layout, BLOCK)
    ty = op(tx, scale=scale)

    # torch result
    rx = mask_tensor(x, layout, BLOCK, value=float("-inf"))
    rx = rx[:, :, : (M // BLOCK) * BLOCK, : (M // BLOCK) * BLOCK]

    ry = torch.softmax(rx * scale, -1)
    ry = block_sparsify_tensor(ry, layout, BLOCK)

    # compare
    t.assert_allclose(ry, ty)
//  [python]
def test_softmax(BLOCK, WIDTH, DTYPE):
    # set seed
    torch.random.manual_seed(0)
    Z, H, M, N = 2, 4, WIDTH, WIDTH
    scale = 0.4

    # create inputs
    layout = torch.randint(2, (H, M // BLOCK, N // BLOCK))
    x = torch.randn((Z, H, M, N), dtype=DTYPE, requires_grad=True, device="cuda")

    # triton result
    op = blocksparse_softmax(layout, BLOCK, device=torch.device("cuda"))
    tx = block_sparsify_tensor(x, layout, BLOCK)
    ty = op(tx, scale=scale)

    # torch result
    rx = mask_tensor(x, layout, BLOCK, value=float("-inf"))
    rx = rx[:, :, : (M // BLOCK) * BLOCK, : (M // BLOCK) * BLOCK]

    ry = torch.softmax(rx * scale, -1)
    ry = block_sparsify_tensor(ry, layout, BLOCK)

    # compare
    torch.testing.assert_close(ry, ty)
//  [python]
def test_self_multihead_attn_additive_mask(self):
        grads = torch.randn_like(self.tst_inputs)
        mask = ((torch.randn(self.sequences, self.seq_length) > 0) * -10000.0).half().cuda()

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=mask,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=mask,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        self.ref_inputs.backward(grads)
        self.tst_inputs.backward(grads)

        torch.testing.assert_allclose(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_self_multihead_attn_additive_mask(self):
        grads = torch.randn_like(self.tst_inputs)
        mask = ((torch.randn(self.sequences, self.seq_length) > 0) * -10000.0).half().cuda()

        ref_outputs, _ = self.ref_layer.forward(
            self.ref_inputs,
            self.ref_inputs,
            self.ref_inputs,
            key_padding_mask=mask,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        tst_outputs, _ = self.tst_layer.forward(
            self.tst_inputs,
            self.tst_inputs,
            self.tst_inputs,
            key_padding_mask=mask,
            need_weights=False,
            attn_mask=None,
            is_training=True,
        )

        self.ref_inputs.backward(grads)
        self.tst_inputs.backward(grads)

        torch.testing.assert_close(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def _buffer_param_check(
        self,
        module: nn.Module,
        module_copy: nn.Module,
        device: torch.device,
        device_copy: torch.device,
        recurse: bool = True,
    ) -> None:
        # check all buffer/param under the module is value-identical
        # but device-different with the copied module.
        for (name, buffer), (name_copy, buffer_copy) in zip(
            list(module.named_buffers(recurse=recurse))
            + list(module.named_parameters(recurse=recurse)),
            list(module_copy.named_buffers(recurse=recurse))
            + list(module_copy.named_parameters(recurse=recurse)),
        ):
            self.assertEquals(name, name_copy)
            actual, expected = buffer.detach().cpu(), buffer_copy.detach().cpu()
            rtol, atol = _get_default_rtol_and_atol(actual, expected)
            self.assertEqual(actual.shape, expected.shape)
            self.assertTrue(torch.allclose(actual, expected, rtol=rtol, atol=atol))
            self.assertEquals(buffer.detach().device, device)
            self.assertEquals(buffer_copy.detach().device, device_copy)
//  [python]
def _buffer_param_check(
        self,
        module: nn.Module,
        module_copy: nn.Module,
        device: torch.device,
        device_copy: torch.device,
        recurse: bool = True,
    ) -> None:
        # check all buffer/param under the module is value-identical
        # but device-different with the copied module.
        for (name, buffer), (name_copy, buffer_copy) in zip(
            list(module.named_buffers(recurse=recurse))
            + list(module.named_parameters(recurse=recurse)),
            list(module_copy.named_buffers(recurse=recurse))
            + list(module_copy.named_parameters(recurse=recurse)),
        ):
            self.assertEquals(name, name_copy)
            actual, expected = buffer.detach().cpu(), buffer_copy.detach().cpu()
            rtol, atol = _get_default_rtol_and_atol(actual, expected)
            torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol)
            self.assertEquals(buffer.detach().device, device)
            self.assertEquals(buffer_copy.detach().device, device_copy)
//  [python]
def _test_dense_to_jagged(
        self,
        num_jagged_dim: int,
        outer_dense_size: int,
        inner_dense_size: int,
        dtype: torch.dtype,
        device: torch.device,
        precompute_total_L: bool,
    ) -> None:
        # Generate multi-dim jagged tensor
        values_2d, offsets, max_lengths = generate_jagged_tensor(
            num_jagged_dim, outer_dense_size, inner_dense_size, dtype, device
        )
        values_2d = values_2d.clone().detach().requires_grad_(True)

        # jagged -> dense
        dense = torch.ops.fbgemm.jagged_to_padded_dense(values_2d, offsets, max_lengths)

        # dense -> jagged (op which is being tested)
        if precompute_total_L:
            total_L = values_2d.size(0)
            jagged_values, jagged_offsets = torch.ops.fbgemm.dense_to_jagged(
                dense, offsets, total_L
            )
        else:
            jagged_values, jagged_offsets = torch.ops.fbgemm.dense_to_jagged(
                dense, offsets
            )

        # jagged -> dense
        dense2 = torch.ops.fbgemm.jagged_to_padded_dense(
            jagged_values, jagged_offsets, max_lengths
        )

        # verify forward
        torch.testing.assert_allclose(dense, dense2)

        # verify backward
        dense.retain_grad()
        ref_output_values = jagged_values.clone().detach().requires_grad_(True)
        ref_values = dense.clone().detach().requires_grad_(True)
        jagged_values.backward(ref_output_values)
        torch.testing.assert_allclose(dense.grad, ref_values)
//  [python]
def _test_dense_to_jagged(
        self,
        num_jagged_dim: int,
        outer_dense_size: int,
        inner_dense_size: int,
        dtype: torch.dtype,
        device: torch.device,
        precompute_total_L: bool,
    ) -> None:
        # Generate multi-dim jagged tensor
        values_2d, offsets, max_lengths = generate_jagged_tensor(
            num_jagged_dim, outer_dense_size, inner_dense_size, dtype, device
        )
        values_2d = values_2d.clone().detach().requires_grad_(True)

        # jagged -> dense
        dense = torch.ops.fbgemm.jagged_to_padded_dense(values_2d, offsets, max_lengths)

        # dense -> jagged (op which is being tested)
        if precompute_total_L:
            total_L = values_2d.size(0)
            jagged_values, jagged_offsets = torch.ops.fbgemm.dense_to_jagged(
                dense, offsets, total_L
            )
        else:
            jagged_values, jagged_offsets = torch.ops.fbgemm.dense_to_jagged(
                dense, offsets
            )

        # jagged -> dense
        dense2 = torch.ops.fbgemm.jagged_to_padded_dense(
            jagged_values, jagged_offsets, max_lengths
        )

        # verify forward
        torch.testing.assert_close(dense, dense2)

        # verify backward
        dense.retain_grad()
        ref_output_values = jagged_values.clone().detach().requires_grad_(True)
        ref_values = dense.clone().detach().requires_grad_(True)
        jagged_values.backward(ref_output_values)
        torch.testing.assert_close(dense.grad, ref_values)
//  [python]
def test_weights(self) -> None:
        batch = dummy_batch()
        self.assertTrue(torch.allclose(batch.weight.sum(), torch.tensor(1.0)))
        a, b = batch.split(1)
        self.assertTrue(torch.allclose(a.weight.sum() + b.weight.sum(), torch.tensor(1.0)))