jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_ms_pair_batch(self):
        x = torch.tensor([1])
        pos = x
        x_target = torch.tensor([2])
        pos_target = x
        ms = [Data(x=x, pos=pos), Data(x=4 * x, pos=4 * pos)]
        ms_target = [Data(x=x_target, pos=pos_target),
                     Data(x=4 * x_target, pos=4 * pos_target)]
        data1 = MultiScalePair(x=x, pos=pos, multiscale=ms,
                               x_target=x_target, pos_target=pos_target,
                               multiscale_target=ms_target)

        x = torch.tensor([3])
        pos = x
        x_target = torch.tensor([4])
        pos_target = x
        ms = [Data(x=x, pos=pos), Data(x=4 * x, pos=4 * pos)]
        ms_target = [Data(x=x_target, pos=pos_target),
                     Data(x=4 * x_target, pos=4 * pos_target)]
        data2 = MultiScalePair(x=x, pos=pos, multiscale=ms,
                               x_target=x_target, pos_target=pos_target,
                               multiscale_target=ms_target)

        batch = PairMultiScaleBatch.from_data_list([data1, data2])
        tt.assert_close(batch.x, torch.tensor([1, 3]))
        tt.assert_close(batch.x_target, torch.tensor([2, 4]))
        tt.assert_close(batch.batch, torch.tensor([0, 1]))
        tt.assert_close(batch.batch_target, torch.tensor([0, 1]))

        ms_batches = batch.multiscale
        tt.assert_close(ms_batches[0].batch, torch.tensor([0, 1]))
        tt.assert_close(ms_batches[1].batch, torch.tensor([0, 1]))
        tt.assert_close(ms_batches[1].x, torch.tensor([4, 12]))

        ms_batches = batch.multiscale_target
        tt.assert_close(ms_batches[0].batch, torch.tensor([0, 1]))
        tt.assert_close(ms_batches[1].batch, torch.tensor([0, 1]))
        tt.assert_close(ms_batches[1].x, torch.tensor([8, 16]))
//  [python]
def test_sparse_to_dense_mask(self, dim_1, dim_2, dim_3):
        indices = np.array([i + 1 for i in range(dim_1)]).astype(np.int32)
        values = np.random.rand(dim_1, dim_2, dim_3).astype(np.float32)
        default_value = np.zeros((dim_2, dim_3)).astype(np.float32)
        mask = [2, 4, 9]

        def sparse_to_dense_mask_ref(return_presence_mask=False):
            ref_op = core.CreateOperator(
                "SparseToDenseMask",
                ["indices", "values", "default_value"],
                ["output", "presence_mask"],
                mask=mask,
                return_presence_mask=return_presence_mask,
            )
            workspace.FeedBlob("indices", indices)
            workspace.FeedBlob("values", values)
            workspace.FeedBlob("default_value", default_value)
            workspace.RunOperatorOnce(ref_op)

            if return_presence_mask:
                return (
                    workspace.FetchBlob("output"),
                    workspace.FetchBlob("presence_mask"),
                )

            return workspace.FetchBlob("output")

        # Testing return_presence_mask = False
        output = sparse_to_dense_mask_ref()
        output = torch.tensor(output)

        a, _ = torch.ops._caffe2.SparseToDenseMask(
            torch.tensor(indices),
            torch.tensor(values),
            torch.tensor(default_value),
            None,
            mask=mask,
        )

        torch.testing.assert_allclose(output, a)

        # Testing return_presence_mask = True
        output, presence_mask = sparse_to_dense_mask_ref(return_presence_mask=True)
        output = torch.tensor(output)
        presence_mask = torch.tensor(presence_mask)

        a, b = torch.ops._caffe2.SparseToDenseMask(
            torch.tensor(indices),
            torch.tensor(values),
            torch.tensor(default_value),
            None,
            mask=mask,
            return_presence_mask=True,
        )

        torch.testing.assert_allclose(output, a)
        torch.testing.assert_allclose(presence_mask, b)
//  [python]
def test_sparse_to_dense_mask(self, dim_1, dim_2, dim_3):
        indices = np.array([i + 1 for i in range(dim_1)]).astype(np.int32)
        values = np.random.rand(dim_1, dim_2, dim_3).astype(np.float32)
        default_value = np.zeros((dim_2, dim_3)).astype(np.float32)
        mask = [2, 4, 9]

        def sparse_to_dense_mask_ref(return_presence_mask=False):
            ref_op = core.CreateOperator(
                "SparseToDenseMask",
                ["indices", "values", "default_value"],
                ["output", "presence_mask"],
                mask=mask,
                return_presence_mask=return_presence_mask,
            )
            workspace.FeedBlob("indices", indices)
            workspace.FeedBlob("values", values)
            workspace.FeedBlob("default_value", default_value)
            workspace.RunOperatorOnce(ref_op)

            if return_presence_mask:
                return (
                    workspace.FetchBlob("output"),
                    workspace.FetchBlob("presence_mask"),
                )

            return workspace.FetchBlob("output")

        # Testing return_presence_mask = False
        output = sparse_to_dense_mask_ref()
        output = torch.tensor(output)

        a, _ = torch.ops._caffe2.SparseToDenseMask(
            torch.tensor(indices),
            torch.tensor(values),
            torch.tensor(default_value),
            None,
            mask=mask,
        )

        torch.testing.assert_close(output, a)

        # Testing return_presence_mask = True
        output, presence_mask = sparse_to_dense_mask_ref(return_presence_mask=True)
        output = torch.tensor(output)
        presence_mask = torch.tensor(presence_mask)

        a, b = torch.ops._caffe2.SparseToDenseMask(
            torch.tensor(indices),
            torch.tensor(values),
            torch.tensor(default_value),
            None,
            mask=mask,
            return_presence_mask=True,
        )

        torch.testing.assert_close(output, a)
        torch.testing.assert_close(presence_mask, b)
//  [python]
def _test_resize_nearest_op(self, device):
        data = np.random.rand(1, 2, 3, 4).astype(np.float32)

        def _resize_nearest_ref(X):
            ref_op = core.CreateOperator(
                "ResizeNearest",
                ["X"],
                ["Y"],
                width_scale=2.0,
                height_scale=1.5,
                order="NCHW",
            )
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _resize_nearest_ref(data)
        actual_output = torch.ops._caffe2.ResizeNearest(
            torch.tensor(data).to(device),
            order="NCHW",
            width_scale=2.0,
            height_scale=1.5,
        )

        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def _test_resize_nearest_op(self, device):
        data = np.random.rand(1, 2, 3, 4).astype(np.float32)

        def _resize_nearest_ref(X):
            ref_op = core.CreateOperator(
                "ResizeNearest",
                ["X"],
                ["Y"],
                width_scale=2.0,
                height_scale=1.5,
                order="NCHW",
            )
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _resize_nearest_ref(data)
        actual_output = torch.ops._caffe2.ResizeNearest(
            torch.tensor(data).to(device),
            order="NCHW",
            width_scale=2.0,
            height_scale=1.5,
        )

        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def _test_roi_align_rotated(self, N, C, H, W, device):
        def rand_rotated_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * 360 - 180,
                ]
            ).astype(np.float32)

        feature = np.random.randn(N, C, H, W).astype(np.float32)
        rois = np.array([rand_rotated_roi() for _ in range(10)])

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlignRotated",
                ["feature", "rois"],
                ["roi_feature"],
                spatial_scale=1.0,
                pooled_h=3,
                pooled_w=3,
                sampling_ratio=0,
            )
            workspace.FeedBlob("feature", _feature)
            workspace.FeedBlob("rois", _rois)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("roi_feature")

        roi_feature_ref = roi_align_ref(feature, rois)
        roi_feature = torch.ops._caffe2.RoIAlignRotated(
            torch.tensor(feature).to(device),
            torch.tensor(rois).to(device),
            order="NCHW",
            spatial_scale=1.0,
            pooled_h=3,
            pooled_w=3,
            sampling_ratio=0,
            aligned=False,
        )
        torch.testing.assert_allclose(roi_feature_ref, roi_feature.cpu())
//  [python]
def _test_roi_align_rotated(self, N, C, H, W, device):
        def rand_rotated_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * 360 - 180,
                ]
            ).astype(np.float32)

        feature = np.random.randn(N, C, H, W).astype(np.float32)
        rois = np.array([rand_rotated_roi() for _ in range(10)])

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlignRotated",
                ["feature", "rois"],
                ["roi_feature"],
                spatial_scale=1.0,
                pooled_h=3,
                pooled_w=3,
                sampling_ratio=0,
            )
            workspace.FeedBlob("feature", _feature)
            workspace.FeedBlob("rois", _rois)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("roi_feature")

        roi_feature_ref = roi_align_ref(feature, rois)
        roi_feature = torch.ops._caffe2.RoIAlignRotated(
            torch.tensor(feature).to(device),
            torch.tensor(rois).to(device),
            order="NCHW",
            spatial_scale=1.0,
            pooled_h=3,
            pooled_w=3,
            sampling_ratio=0,
            aligned=False,
        )
        torch.testing.assert_close(roi_feature_ref, roi_feature.cpu())
//  [python]
def test_copy_between_cpu_and_gpu(self):
        x_cpu_ref = torch.tensor([1., 2., 3.])
        x_gpu_ref = x_cpu_ref.to("cuda")

        x_gpu = torch.ops._caffe2.CopyCPUToGPU(x_cpu_ref)
        torch.testing.assert_allclose(x_gpu, x_gpu_ref)
        x_cpu = torch.ops._caffe2.CopyGPUToCPU(x_gpu)
        torch.testing.assert_allclose(x_cpu, x_cpu_ref)
//  [python]
def test_copy_between_cpu_and_gpu(self):
        x_cpu_ref = torch.tensor([1., 2., 3.])
        x_gpu_ref = x_cpu_ref.to("cuda")

        x_gpu = torch.ops._caffe2.CopyCPUToGPU(x_cpu_ref)
        torch.testing.assert_close(x_gpu, x_gpu_ref)
        x_cpu = torch.ops._caffe2.CopyGPUToCPU(x_gpu)
        torch.testing.assert_close(x_cpu, x_cpu_ref)
//  [python]
def test_alias_with_name_is_in_place(self):
        device = "cuda" if workspace.has_cuda_support else "cpu"
        x = torch.tensor([3., 42.]).to(device=device)
        y = torch.ops._caffe2.AliasWithName(x, "new_name")
        x[1] = 6
        torch.testing.assert_allclose(x, torch.tensor([3., 6.]).to(device=device))
        # y should also change because y is alias of x
        torch.testing.assert_allclose(y, torch.tensor([3., 6.]).to(device=device))
//  [python]
def test_alias_with_name_is_in_place(self):
        device = "cuda" if workspace.has_cuda_support else "cpu"
        x = torch.tensor([3., 42.]).to(device=device)
        y = torch.ops._caffe2.AliasWithName(x, "new_name")
        x[1] = 6
        torch.testing.assert_close(x, torch.tensor([3., 6.]).to(device=device))
        # y should also change because y is alias of x
        torch.testing.assert_close(y, torch.tensor([3., 6.]).to(device=device))
//  [python]
def test_piecewise_linear_op(self, binary_input):
        if binary_input:
            num_dims = 1
        else:
            num_dims = 3
        data = np.random.rand(1024, num_dims).astype(np.float32)
        slopes = np.zeros(4 * num_dims).astype(np.float32)
        bounds = np.sort(
            np.random.rand(5, num_dims).astype(np.float32), axis=0
        ).flatten("F")
        intercepts = np.random.rand(4 * num_dims).astype(np.float32)

        def _piecewise_linear_ref(X):
            ref_op = core.CreateOperator(
                "PiecewiseLinearTransform",
                ["data", "bounds", "slopes", "intercepts"],
                ["calibrated"],
                binary=binary_input,
            )
            workspace.FeedBlob("data", X)
            workspace.FeedBlob("bounds", bounds)
            workspace.FeedBlob("slopes", slopes)
            workspace.FeedBlob("intercepts", intercepts)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("calibrated")

        expected_output = _piecewise_linear_ref(data)
        actual_output = torch.ops._caffe2.PiecewiseLinearTransform(
            torch.tensor(data),
            bounds.tolist(),
            slopes.tolist(),
            intercepts.tolist(),
            binary_input,
        )

        torch.testing.assert_allclose(torch.tensor(expected_output), actual_output)
//  [python]
def test_piecewise_linear_op(self, binary_input):
        if binary_input:
            num_dims = 1
        else:
            num_dims = 3
        data = np.random.rand(1024, num_dims).astype(np.float32)
        slopes = np.zeros(4 * num_dims).astype(np.float32)
        bounds = np.sort(
            np.random.rand(5, num_dims).astype(np.float32), axis=0
        ).flatten("F")
        intercepts = np.random.rand(4 * num_dims).astype(np.float32)

        def _piecewise_linear_ref(X):
            ref_op = core.CreateOperator(
                "PiecewiseLinearTransform",
                ["data", "bounds", "slopes", "intercepts"],
                ["calibrated"],
                binary=binary_input,
            )
            workspace.FeedBlob("data", X)
            workspace.FeedBlob("bounds", bounds)
            workspace.FeedBlob("slopes", slopes)
            workspace.FeedBlob("intercepts", intercepts)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("calibrated")

        expected_output = _piecewise_linear_ref(data)
        actual_output = torch.ops._caffe2.PiecewiseLinearTransform(
            torch.tensor(data),
            bounds.tolist(),
            slopes.tolist(),
            intercepts.tolist(),
            binary_input,
        )

        torch.testing.assert_close(torch.tensor(expected_output), actual_output)
//  [python]
def test_index_hash_op(self):
        data = np.random.randint(low=0, high=1000, size=(4, 4, 4))

        def _index_hash_ref(X):
            ref_op = core.CreateOperator("IndexHash", ["X"], ["Y"], seed=0, modulo=100)
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _index_hash_ref(data)
        actual_output = torch.ops._caffe2.IndexHash(
            torch.tensor(data), seed=0, modulo=100
        )

        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_index_hash_op(self):
        data = np.random.randint(low=0, high=1000, size=(4, 4, 4))

        def _index_hash_ref(X):
            ref_op = core.CreateOperator("IndexHash", ["X"], ["Y"], seed=0, modulo=100)
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _index_hash_ref(data)
        actual_output = torch.ops._caffe2.IndexHash(
            torch.tensor(data), seed=0, modulo=100
        )

        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_bucketize_op(self):
        data = np.random.rand(8, 10).astype(np.float32) * 1000
        boundaries = np.array([1, 10, 100, 1000, 100000]).astype(np.float32)

        def _bucketize_ref(X):
            ref_op = core.CreateOperator(
                "Bucketize", ["X"], ["Y"], boundaries=boundaries
            )
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _bucketize_ref(data)
        actual_output = torch.ops._caffe2.Bucketize(torch.tensor(data), boundaries)
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_bucketize_op(self):
        data = np.random.rand(8, 10).astype(np.float32) * 1000
        boundaries = np.array([1, 10, 100, 1000, 100000]).astype(np.float32)

        def _bucketize_ref(X):
            ref_op = core.CreateOperator(
                "Bucketize", ["X"], ["Y"], boundaries=boundaries
            )
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _bucketize_ref(data)
        actual_output = torch.ops._caffe2.Bucketize(torch.tensor(data), boundaries)
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_logit(self, X, eps):
        def ref(X, eps):
            ref_op = core.CreateOperator("Logit", ["X"], ["Y"], eps=eps)
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = ref(X, eps)
        actual_output = torch.ops._caffe2.Logit(torch.tensor(X), eps)
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_logit(self, X, eps):
        def ref(X, eps):
            ref_op = core.CreateOperator("Logit", ["X"], ["Y"], eps=eps)
            workspace.FeedBlob("X", X)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = ref(X, eps)
        actual_output = torch.ops._caffe2.Logit(torch.tensor(X), eps)
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_pack_segments(self):
        s = torch.rand(3, 3, 3)
        lengths = torch.tensor([2, 1])
        packed_tensor, _ = torch.ops._caffe2.PackSegments(lengths, s)
        self.assertEqual(packed_tensor.numpy().shape, (2, 2, 3, 3))
        unpacked_tensor = torch.ops._caffe2.UnpackSegments(lengths, packed_tensor)
        torch.testing.assert_allclose(s, unpacked_tensor)
//  [python]
def test_pack_segments(self):
        s = torch.rand(3, 3, 3)
        lengths = torch.tensor([2, 1])
        packed_tensor, _ = torch.ops._caffe2.PackSegments(lengths, s)
        self.assertEqual(packed_tensor.numpy().shape, (2, 2, 3, 3))
        unpacked_tensor = torch.ops._caffe2.UnpackSegments(lengths, packed_tensor)
        torch.testing.assert_close(s, unpacked_tensor)
//  [python]
def test_save_and_load_checkpoint(self):
        model = torchvision.models.resnet18(pretrained=False)
        utils.save_checkpoint(model, epoch=100, filename='tmp.pth', save_arch=True)

        loaded_model = utils.load_model('tmp.pth')

        torch.testing.assert_allclose(model.conv1.weight, loaded_model.conv1.weight)

        model.conv1.weight = nn.Parameter(torch.zeros_like(model.conv1.weight))
        model = utils.load_checkpoint('tmp.pth', model=model)['model']

        assert (model.conv1.weight != 0).any()
//  [python]
def test_save_and_load_checkpoint(self):
        model = torchvision.models.resnet18(pretrained=False)
        utils.save_checkpoint(model, epoch=100, filename='tmp.pth', save_arch=True)

        loaded_model = utils.load_model('tmp.pth')

        torch.testing.assert_close(model.conv1.weight, loaded_model.conv1.weight)

        model.conv1.weight = nn.Parameter(torch.zeros_like(model.conv1.weight))
        model = utils.load_checkpoint('tmp.pth', model=model)['model']

        assert (model.conv1.weight != 0).any()
//  [python]
def test_percentile(self):
        original_values = np.array([[3.0, 5.0, 3], [5.0, 1.0, 6.0]]).astype(np.float32)
        value_to_pct = np.array([[3, 0.2], [5, 0.5], [1, 0.3], [3, 0.6]]).astype(
            np.float32
        )
        lengths = np.array([2, 1, 1]).astype(np.int32)

        def _percentile_ref(original_values, value_to_pct, lengths):
            ref_op = core.CreateOperator(
                "Percentile", ["original_values", "value_to_pct", "lengths"], ["Y"]
            )
            workspace.FeedBlob("original_values", original_values)
            workspace.FeedBlob("value_to_pct", value_to_pct)
            workspace.FeedBlob("lengths", lengths)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _percentile_ref(original_values, value_to_pct, lengths)
        actual_output = torch.ops._caffe2.Percentile(
            torch.tensor(original_values),
            torch.tensor(value_to_pct),
            torch.tensor(lengths),
        )
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_percentile(self):
        original_values = np.array([[3.0, 5.0, 3], [5.0, 1.0, 6.0]]).astype(np.float32)
        value_to_pct = np.array([[3, 0.2], [5, 0.5], [1, 0.3], [3, 0.6]]).astype(
            np.float32
        )
        lengths = np.array([2, 1, 1]).astype(np.int32)

        def _percentile_ref(original_values, value_to_pct, lengths):
            ref_op = core.CreateOperator(
                "Percentile", ["original_values", "value_to_pct", "lengths"], ["Y"]
            )
            workspace.FeedBlob("original_values", original_values)
            workspace.FeedBlob("value_to_pct", value_to_pct)
            workspace.FeedBlob("lengths", lengths)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _percentile_ref(original_values, value_to_pct, lengths)
        actual_output = torch.ops._caffe2.Percentile(
            torch.tensor(original_values),
            torch.tensor(value_to_pct),
            torch.tensor(lengths),
        )
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_batch_bucket_one_hot_op(self):
        data = np.array([[2, 3], [4, 1], [2, 5]]).astype(np.float32)
        lengths = np.array([2, 3]).astype(np.int32)
        boundaries = np.array([0.1, 2.5, 1, 3.1, 4.5]).astype(np.float32)

        def _batch_bucket_one_hot_ref(data, lengths, boundaries):
            ref_op = core.CreateOperator(
                "BatchBucketOneHot", ["data", "lengths", "boundaries"], ["Y"]
            )
            workspace.FeedBlob("data", data)
            workspace.FeedBlob("lengths", lengths)
            workspace.FeedBlob("boundaries", boundaries)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _batch_bucket_one_hot_ref(data, lengths, boundaries)
        actual_output = torch.ops._caffe2.BatchBucketOneHot(
            torch.tensor(data), torch.tensor(lengths), torch.tensor(boundaries)
        )
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_batch_bucket_one_hot_op(self):
        data = np.array([[2, 3], [4, 1], [2, 5]]).astype(np.float32)
        lengths = np.array([2, 3]).astype(np.int32)
        boundaries = np.array([0.1, 2.5, 1, 3.1, 4.5]).astype(np.float32)

        def _batch_bucket_one_hot_ref(data, lengths, boundaries):
            ref_op = core.CreateOperator(
                "BatchBucketOneHot", ["data", "lengths", "boundaries"], ["Y"]
            )
            workspace.FeedBlob("data", data)
            workspace.FeedBlob("lengths", lengths)
            workspace.FeedBlob("boundaries", boundaries)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _batch_bucket_one_hot_ref(data, lengths, boundaries)
        actual_output = torch.ops._caffe2.BatchBucketOneHot(
            torch.tensor(data), torch.tensor(lengths), torch.tensor(boundaries)
        )
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_smoothness(self):
        core_smth = C.Smoothness()
        torch_smth = T.Smoothness()

        for _ in range(3):
            z = np.repeat(
                np.repeat(np.arange(16)[None, None, :], 16, axis=0), 8, axis=1
            )
            a = np.random.randn(16, 3, 16)

            ztm = torch.from_numpy(z)
            atm = torch.from_numpy(a)

            core_smth.update_state(z, a)
            torch_smth.update(ztm, atm)

        val = core_smth.compute()
        valtm = torch_smth.compute()

        np.testing.assert_allclose(val, valtm)

        torch.testing.assert_allclose(val, valtm)
//  [python]
def test_smoothness(self):
        core_smth = C.Smoothness()
        torch_smth = T.Smoothness()

        for _ in range(3):
            z = np.repeat(
                np.repeat(np.arange(16)[None, None, :], 16, axis=0), 8, axis=1
            )
            a = np.random.randn(16, 3, 16)

            ztm = torch.from_numpy(z)
            atm = torch.from_numpy(a)

            core_smth.update_state(z, a)
            torch_smth.update(ztm, atm)

        val = core_smth.compute()
        valtm = torch_smth.compute()

        np.testing.assert_allclose(val, valtm)

        torch.testing.assert_close(val, valtm)
//  [python]
def test_batch_symeig_backward():
    input = torch.randn(8*64*32, 4, 4).float()
    input1 = input.clone()
    input.requires_grad = True
    w, v = S.symeig(input)
    (w.sum() + v.sum()).backward()
    # print(input.grad)
    for i in range(input1.size(0)):
        in1 = input1[i]
        in1.requires_grad = True
        wi, vi = S.symeig(in1)
        (wi.sum() + vi.sum()).backward()
        # print(in1.grad)
        torch.testing.assert_allclose(input.grad[i], in1.grad)
//  [python]
def test_batch_symeig_backward():
    input = torch.randn(8*64*32, 4, 4).float()
    input1 = input.clone()
    input.requires_grad = True
    w, v = S.symeig(input)
    (w.sum() + v.sum()).backward()
    # print(input.grad)
    for i in range(input1.size(0)):
        in1 = input1[i]
        in1.requires_grad = True
        wi, vi = S.symeig(in1)
        (wi.sum() + vi.sum()).backward()
        # print(in1.grad)
        torch.testing.assert_close(input.grad[i], in1.grad)
//  [python]
def test_indexedconv11(self):
        torch.testing.assert_allclose(self.conv11(self.data_1),
                                      torch.tensor([8, -2, 8, -2, 5, -2, 8, -2, 8], dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_indexedconv11(self):
        torch.testing.assert_close(self.conv11(self.data_1),
                                      torch.tensor([8, -2, 8, -2, 5, -2, 8, -2, 8], dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_maxpool1(self):
        torch.testing.assert_allclose(self.maxpool(self.data_1),
                                      torch.tensor([[[4.,  5.,  4.,  5.]]]))
//  [python]
def test_maxpool1(self):
        torch.testing.assert_close(self.maxpool(self.data_1),
                                      torch.tensor([[[4.,  5.,  4.,  5.]]]))
//  [python]
def test_indexedconv13(self):
        torch.testing.assert_allclose(self.conv13(self.data_1),
                                      torch.tensor([[8, -2, 8, -2, 5, -2, 8, -2, 8],
                                                    [8, -2, 8, -2, 5, -2, 8, -2, 8],
                                                    [8, -2, 8, -2, 5, -2, 8, -2, 8]], dtype=torch.float).unsqueeze(0))
//  [python]
def test_indexedconv13(self):
        torch.testing.assert_close(self.conv13(self.data_1),
                                      torch.tensor([[8, -2, 8, -2, 5, -2, 8, -2, 8],
                                                    [8, -2, 8, -2, 5, -2, 8, -2, 8],
                                                    [8, -2, 8, -2, 5, -2, 8, -2, 8]], dtype=torch.float).unsqueeze(0))
//  [python]
def test_merge_id_lists(self, lengths_0, lengths_1):
        def _merge_id_lists(lengths, values):
            ref_op = core.CreateOperator(
                "MergeIdLists",
                ["lengths_0", "values_0", "lengths_1", "values_1"],
                ["merged_lengths", "merged_values"],
            )
            workspace.FeedBlob("lengths_0", lengths[0])
            workspace.FeedBlob("values_0", values[0])
            workspace.FeedBlob("lengths_1", lengths[1])
            workspace.FeedBlob("values_1", values[1])
            workspace.RunOperatorOnce(ref_op)
            return (
                workspace.FetchBlob("merged_lengths"),
                workspace.FetchBlob("merged_values"),
            )

        lengths = [
            np.array([lengths_0]).astype(np.int32),
            np.array([lengths_1]).astype(np.int32),
        ]
        values = [
            np.random.choice(np.arange(0, 10), size=lengths_0, replace=False).astype(
                np.int32
            ),
            np.random.choice(np.arange(10, 20), size=lengths_1, replace=False).astype(
                np.int32
            ),
        ]

        expected_merged_lengths, expected_merged_values = _merge_id_lists(
            lengths, values
        )
        output_merged_lengths, output_merged_values = torch.ops._caffe2.MergeIdLists(
            [
                torch.tensor(lengths[0]),
                torch.tensor(values[0]),
                torch.tensor(lengths[1]),
                torch.tensor(values[1]),
            ]
        )
        torch.testing.assert_allclose(expected_merged_lengths, output_merged_lengths)
        torch.testing.assert_allclose(expected_merged_values, output_merged_values)
//  [python]
def test_merge_id_lists(self, lengths_0, lengths_1):
        def _merge_id_lists(lengths, values):
            ref_op = core.CreateOperator(
                "MergeIdLists",
                ["lengths_0", "values_0", "lengths_1", "values_1"],
                ["merged_lengths", "merged_values"],
            )
            workspace.FeedBlob("lengths_0", lengths[0])
            workspace.FeedBlob("values_0", values[0])
            workspace.FeedBlob("lengths_1", lengths[1])
            workspace.FeedBlob("values_1", values[1])
            workspace.RunOperatorOnce(ref_op)
            return (
                workspace.FetchBlob("merged_lengths"),
                workspace.FetchBlob("merged_values"),
            )

        lengths = [
            np.array([lengths_0]).astype(np.int32),
            np.array([lengths_1]).astype(np.int32),
        ]
        values = [
            np.random.choice(np.arange(0, 10), size=lengths_0, replace=False).astype(
                np.int32
            ),
            np.random.choice(np.arange(10, 20), size=lengths_1, replace=False).astype(
                np.int32
            ),
        ]

        expected_merged_lengths, expected_merged_values = _merge_id_lists(
            lengths, values
        )
        output_merged_lengths, output_merged_values = torch.ops._caffe2.MergeIdLists(
            [
                torch.tensor(lengths[0]),
                torch.tensor(values[0]),
                torch.tensor(lengths[1]),
                torch.tensor(values[1]),
            ]
        )
        torch.testing.assert_close(expected_merged_lengths, output_merged_lengths)
        torch.testing.assert_close(expected_merged_values, output_merged_values)
//  [python]
def test_indexedconv31(self):
        torch.testing.assert_allclose(self.conv31(self.data_3),
                                      torch.tensor([22, -8, 22, -8, 13, -8, 22, -8, 22], dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_indexedconv31(self):
        torch.testing.assert_close(self.conv31(self.data_3),
                                      torch.tensor([22, -8, 22, -8, 13, -8, 22, -8, 22], dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_averagepool1(self):
        torch.testing.assert_allclose(self.avgpool(self.data_1),
                                      torch.tensor([[[1.25,  1.75,  1.5,  2.25]]]))
//  [python]
def test_averagepool1(self):
        torch.testing.assert_close(self.avgpool(self.data_1),
                                      torch.tensor([[[1.25,  1.75,  1.5,  2.25]]]))
//  [python]
def test_maxpool3(self):
        torch.testing.assert_allclose(self.maxpool(self.data_3),
                                      torch.tensor([[[4., 5., 4., 5.],
                                                     [4., 5., 4., 5.],
                                                     [4., 5., 4., 5.]]]))
//  [python]
def test_maxpool3(self):
        torch.testing.assert_close(self.maxpool(self.data_3),
                                      torch.tensor([[[4., 5., 4., 5.],
                                                     [4., 5., 4., 5.],
                                                     [4., 5., 4., 5.]]]))
//  [python]
def test_averagepool3(self):
        torch.testing.assert_allclose(self.avgpool(self.data_3),
                                      torch.tensor([[[1.25,  1.75,  1.5,  2.25],
                                                     [1.25,  1.75,  1.5,  2.25],
                                                     [1.25,  1.75,  1.5,  2.25]]]))
//  [python]
def test_averagepool3(self):
        torch.testing.assert_close(self.avgpool(self.data_3),
                                      torch.tensor([[[1.25,  1.75,  1.5,  2.25],
                                                     [1.25,  1.75,  1.5,  2.25],
                                                     [1.25,  1.75,  1.5,  2.25]]]))
//  [python]
def test_matting():
    # binary (ceiling the index)
    torch.testing.assert_allclose(
        matting(torch.tensor(0.), 4, binary=True).squeeze(),
        torch.tensor([1., 0., 0., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(1 / 4), 4, binary=True).squeeze(),
        torch.tensor([1., 0., 0., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(2 / 4), 4, binary=True).squeeze(),
        torch.tensor([0., 1., 0., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(3 / 4), 4, binary=True).squeeze(),
        torch.tensor([0., 0., 1., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(1.), 4, binary=True).squeeze(),
        torch.tensor([0., 0., 0., 1.])
    )

    # non-binary
    torch.testing.assert_allclose(
        matting(torch.tensor(0.), 4, binary=False).squeeze(),
        torch.tensor([1., 0., 0., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(1 / 4), 4, binary=False).squeeze(),
        torch.tensor([1., 1., 0., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(2 / 4), 4, binary=False).squeeze(),
        torch.tensor([0., 1., 1., 0.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(3 / 4), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 1., 1.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(1.), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 0., 1.])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(1. - 1 / 4 * 0.1), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 0.1, 1.0])
    )
    torch.testing.assert_allclose(
        matting(torch.tensor(1. - 1 / 4 * 0.9), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 0.9, 1.0])
    )
//  [python]
def test_matting():
    # binary (ceiling the index)
    torch.testing.assert_close(
        matting(torch.tensor(0.), 4, binary=True).squeeze(),
        torch.tensor([1., 0., 0., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(1 / 4), 4, binary=True).squeeze(),
        torch.tensor([1., 0., 0., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(2 / 4), 4, binary=True).squeeze(),
        torch.tensor([0., 1., 0., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(3 / 4), 4, binary=True).squeeze(),
        torch.tensor([0., 0., 1., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(1.), 4, binary=True).squeeze(),
        torch.tensor([0., 0., 0., 1.])
    )

    # non-binary
    torch.testing.assert_close(
        matting(torch.tensor(0.), 4, binary=False).squeeze(),
        torch.tensor([1., 0., 0., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(1 / 4), 4, binary=False).squeeze(),
        torch.tensor([1., 1., 0., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(2 / 4), 4, binary=False).squeeze(),
        torch.tensor([0., 1., 1., 0.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(3 / 4), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 1., 1.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(1.), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 0., 1.])
    )
    torch.testing.assert_close(
        matting(torch.tensor(1. - 1 / 4 * 0.1), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 0.1, 1.0])
    )
    torch.testing.assert_close(
        matting(torch.tensor(1. - 1 / 4 * 0.9), 4, binary=False).squeeze(),
        torch.tensor([0., 0., 0.9, 1.0])
    )
//  [python]
def test_simplecpu(self):
        points = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[-1, 1, 0], [0, 0, 10], [0, 0, 2]]]).float()
        idx = fps(points, 2, False)
        torch.testing.assert_allclose(idx, torch.tensor([[0, 2], [0, 1]]))
//  [python]
def test_simplecpu(self):
        points = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[-1, 1, 0], [0, 0, 10], [0, 0, 2]]]).float()
        idx = fps(points, 2, False)
        torch.testing.assert_close(idx, torch.tensor([[0, 2], [0, 1]]))
//  [python]
def test_approximate_derivatives_tensor_functionality(self) -> None:
        """
        Tests the numerical accuracy of approximate_derivatives_tensor.
        """
        input_y = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.float64)
        input_x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)
        window_length = 3  # Currently the only supported value
        poly_order = 2
        deriv_order = 1

        expected_output = torch.tensor([[1, 1, 1, 1, 1]], dtype=torch.float64)
        actual_output = approximate_derivatives_tensor(input_y, input_x, window_length, poly_order, deriv_order)

        torch.testing.assert_allclose(expected_output, actual_output)
//  [python]
def test_approximate_derivatives_tensor_functionality(self) -> None:
        """
        Tests the numerical accuracy of approximate_derivatives_tensor.
        """
        input_y = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.float64)
        input_x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)
        window_length = 3  # Currently the only supported value
        poly_order = 2
        deriv_order = 1

        expected_output = torch.tensor([[1, 1, 1, 1, 1]], dtype=torch.float64)
        actual_output = approximate_derivatives_tensor(input_y, input_x, window_length, poly_order, deriv_order)

        torch.testing.assert_close(expected_output, actual_output)
//  [python]
def test_unwrap_scripts_properly(self) -> None:
        """
        Tests that unwrap scripts properly.
        """

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

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                output = unwrap(x, dim=-1)

                return output

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

        test_x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)

        py_result = to_script.forward(test_x)
        script_result = scripted.forward(test_x)

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

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

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                output = unwrap(x, dim=-1)

                return output

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

        test_x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)

        py_result = to_script.forward(test_x)
        script_result = scripted.forward(test_x)

        torch.testing.assert_close(py_result, script_result)
//  [python]
def _run_correctness_test(self, opt, model, override=None):
        for i in range(10):
            fmodel = higher.patch.monkeypatch(model)
            diffopt = higher.optim.get_diff_optim(
                opt, model.parameters(), fmodel, override=override
            )
            for j in range(3):
                opt.zero_grad()
                x = torch.rand(10, 4)

                y_model = model(x)
                y_fmodel = fmodel(x)

                loss_model = y_model.pow(2).sum()
                loss_fmodel = y_fmodel.pow(2).sum()

                loss_model.backward()
                diffopt.step(loss_fmodel)
                opt.step()
                self.assertEqual(
                    len(list(model.parameters())),
                    len(list(fmodel.parameters()))
                )

            for p, fp in zip(model.parameters(), fmodel.parameters()):
                torch.testing.assert_allclose(p, fp, atol=1e-5, rtol=1e-1)
//  [python]
def _run_correctness_test(self, opt, model, override=None):
        for i in range(10):
            fmodel = higher.patch.monkeypatch(model)
            diffopt = higher.optim.get_diff_optim(
                opt, model.parameters(), fmodel, override=override
            )
            for j in range(3):
                opt.zero_grad()
                x = torch.rand(10, 4)

                y_model = model(x)
                y_fmodel = fmodel(x)

                loss_model = y_model.pow(2).sum()
                loss_fmodel = y_fmodel.pow(2).sum()

                loss_model.backward()
                diffopt.step(loss_fmodel)
                opt.step()
                self.assertEqual(
                    len(list(model.parameters())),
                    len(list(fmodel.parameters()))
                )

            for p, fp in zip(model.parameters(), fmodel.parameters()):
                torch.testing.assert_close(p, fp, atol=1e-5, rtol=1e-1)
//  [python]
def test_approximate_derivatives_tensor_scripts_properly(self) -> None:
        """
        Tests that approximate_derivatives_tensor scripts properly.
        """

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

            def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
                output = approximate_derivatives_tensor(y, x, window_length=3, poly_order=2, deriv_order=1)

                return output

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

        test_y = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.float64)
        test_x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)

        py_result = to_script.forward(test_x, test_y)
        script_result = scripted.forward(test_x, test_y)

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

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

            def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
                output = approximate_derivatives_tensor(y, x, window_length=3, poly_order=2, deriv_order=1)

                return output

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

        test_y = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.float64)
        test_x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)

        py_result = to_script.forward(test_x, test_y)
        script_result = scripted.forward(test_x, test_y)

        torch.testing.assert_close(py_result, script_result)
//  [python]
def test_int_mm(self, device, dtype):
        from torchao.kernel import intmm

        dtype = torch.bfloat16
        m, k, n = (128, 64, 16)
        x = torch.randn(m, k, dtype=dtype, device=device)
        w = torch.randn(n, k, dtype=dtype, device=device).t()
        x_int = x.to(dtype=torch.int8)
        w_int = w.to(dtype=torch.int8)
        out32_1 = intmm.safe_int_mm(x_int, w_int)
        assert out32_1.dtype == torch.int32
        out32_2 = intmm.int_matmul(x_int, w_int)
        assert out32_2.dtype == out32_1.dtype
        torch.testing.assert_allclose(out32_1, out32_2)
//  [python]
def test_int_mm(self, device, dtype):
        from torchao.kernel import intmm

        dtype = torch.bfloat16
        m, k, n = (128, 64, 16)
        x = torch.randn(m, k, dtype=dtype, device=device)
        w = torch.randn(n, k, dtype=dtype, device=device).t()
        x_int = x.to(dtype=torch.int8)
        w_int = w.to(dtype=torch.int8)
        out32_1 = intmm.safe_int_mm(x_int, w_int)
        assert out32_1.dtype == torch.int32
        out32_2 = intmm.int_matmul(x_int, w_int)
        assert out32_2.dtype == out32_1.dtype
        torch.testing.assert_close(out32_1, out32_2)
//  [python]
def test_equality(self):
        state, action, reward, next_state, done = self.init()
        o = Observation(state, action, reward, next_state, done).to_torch()
        for x, y in zip(
            o, Observation(state, action, reward, next_state, done).to_torch()
        ):
            torch.testing.assert_allclose(x, y)
        assert o is not Observation(state, action, reward, next_state, done)
//  [python]
def test_equality(self):
        state, action, reward, next_state, done = self.init()
        o = Observation(state, action, reward, next_state, done).to_torch()
        for x, y in zip(
            o, Observation(state, action, reward, next_state, done).to_torch()
        ):
            torch.testing.assert_close(x, y)
        assert o is not Observation(state, action, reward, next_state, done)
//  [python]
def test_forward(self):
        print("\nTesting positional encoding...")
        encoding = PositionalEncoding()
        self.assertIsNotNone(encoding)

        # Create dummy data
        dummy_input = np.broadcast_to(np.linspace(-1, 1, 3)[None, :], (1, 3))
        dummy_input = torch.tensor(dummy_input, dtype=torch.float32)

        file_path = os.path.join(
            os.path.dirname(__file__), "data/dummy_positional_encoding_output.npy"
        )
        dummy_output = np.load(file_path)

        res = encoding(dummy_input)
        torch.testing.assert_allclose(res, dummy_output)

        print("Positional encoding is executed successfully. \n")
//  [python]
def test_forward(self):
        print("\nTesting positional encoding...")
        encoding = PositionalEncoding()
        self.assertIsNotNone(encoding)

        # Create dummy data
        dummy_input = np.broadcast_to(np.linspace(-1, 1, 3)[None, :], (1, 3))
        dummy_input = torch.tensor(dummy_input, dtype=torch.float32)

        file_path = os.path.join(
            os.path.dirname(__file__), "data/dummy_positional_encoding_output.npy"
        )
        dummy_output = np.load(file_path)

        res = encoding(dummy_input)
        torch.testing.assert_close(res, dummy_output)

        print("Positional encoding is executed successfully. \n")