jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_pad_polylines(self) -> None:
        """Test padding polyline features."""
        num_elements = 10
        num_points = 20
        num_features = 2
        input = torch.ones((num_elements, num_points, num_features), dtype=torch.float32)

        elements_padded = pad_polylines(input, 20, 0)
        self.assertEqual(elements_padded.dtype, torch.float32)
        self.assertEqual(elements_padded.shape, (20, num_points, num_features))
        torch.testing.assert_allclose(
            elements_padded[10:, :, :], torch.zeros((10, num_points, num_features), dtype=torch.float32)
        )

        points_padded = pad_polylines(input, 30, 1)
        self.assertEqual(points_padded.dtype, torch.float32)
        self.assertEqual(points_padded.shape, (num_elements, 30, num_features))
        torch.testing.assert_allclose(
            points_padded[:, 20:, :], torch.zeros((num_elements, 10, num_features), dtype=torch.float32)
        )

        features_padded = pad_polylines(input, 3, 2)
        self.assertEqual(features_padded.dtype, torch.float32)
        self.assertEqual(features_padded.shape, (num_elements, num_points, 3))
        torch.testing.assert_allclose(
            features_padded[:, :, 2:], torch.zeros((num_elements, num_points, 1), dtype=torch.float32)
        )
//  [python]
def test_pad_polylines(self) -> None:
        """Test padding polyline features."""
        num_elements = 10
        num_points = 20
        num_features = 2
        input = torch.ones((num_elements, num_points, num_features), dtype=torch.float32)

        elements_padded = pad_polylines(input, 20, 0)
        self.assertEqual(elements_padded.dtype, torch.float32)
        self.assertEqual(elements_padded.shape, (20, num_points, num_features))
        torch.testing.assert_close(
            elements_padded[10:, :, :], torch.zeros((10, num_points, num_features), dtype=torch.float32)
        )

        points_padded = pad_polylines(input, 30, 1)
        self.assertEqual(points_padded.dtype, torch.float32)
        self.assertEqual(points_padded.shape, (num_elements, 30, num_features))
        torch.testing.assert_close(
            points_padded[:, 20:, :], torch.zeros((num_elements, 10, num_features), dtype=torch.float32)
        )

        features_padded = pad_polylines(input, 3, 2)
        self.assertEqual(features_padded.dtype, torch.float32)
        self.assertEqual(features_padded.shape, (num_elements, num_points, 3))
        torch.testing.assert_close(
            features_padded[:, :, 2:], torch.zeros((num_elements, num_points, 1), dtype=torch.float32)
        )
//  [python]
def test_batch(self):
        x = torch.tensor([1])
        pos = x
        d1 = Data(x=x, pos=pos)
        d2 = Data(x=4 * x, pos=4 * pos)
        data1 = MultiScaleData(x=x, pos=pos, multiscale=[d1, d2])

        x = torch.tensor([2])
        pos = x
        d1 = Data(x=x, pos=pos)
        d2 = Data(x=4 * x, pos=4 * pos)
        data2 = MultiScaleData(x=x, pos=pos, multiscale=[d1, d2])

        batch = MultiScaleBatch.from_data_list([data1, data2])
        tt.assert_allclose(batch.x, torch.tensor([1, 2]))
        tt.assert_allclose(batch.batch, torch.tensor([0, 1]))

        ms_batches = batch.multiscale
        tt.assert_allclose(ms_batches[0].batch, torch.tensor([0, 1]))
        tt.assert_allclose(ms_batches[1].batch, torch.tensor([0, 1]))
        tt.assert_allclose(ms_batches[1].x, torch.tensor([4, 8]))
//  [python]
def test_batch(self):
        x = torch.tensor([1])
        pos = x
        d1 = Data(x=x, pos=pos)
        d2 = Data(x=4 * x, pos=4 * pos)
        data1 = MultiScaleData(x=x, pos=pos, multiscale=[d1, d2])

        x = torch.tensor([2])
        pos = x
        d1 = Data(x=x, pos=pos)
        d2 = Data(x=4 * x, pos=4 * pos)
        data2 = MultiScaleData(x=x, pos=pos, multiscale=[d1, d2])

        batch = MultiScaleBatch.from_data_list([data1, data2])
        torch.testing.assert_close(batch.x, torch.tensor([1, 2]))
        torch.testing.assert_close(batch.batch, torch.tensor([0, 1]))

        ms_batches = batch.multiscale
        torch.testing.assert_close(ms_batches[0].batch, torch.tensor([0, 1]))
        torch.testing.assert_close(ms_batches[1].batch, torch.tensor([0, 1]))
        torch.testing.assert_close(ms_batches[1].x, torch.tensor([4, 8]))
//  [python]
def _test_gelu_op(self, X, fast_gelu, device):
        def _gelu_ref(_X):
            return (_X * norm.cdf(_X).astype(np.float32), )
        expected_output, = _gelu_ref(X)
        actual_output = torch.ops._caffe2.Gelu(torch.tensor(X), fast_gelu)

        rtol = 1e-3 if fast_gelu else 1e-4
        atol = 1e-5
        torch.testing.assert_allclose(
            expected_output, actual_output.cpu(), rtol=rtol, atol=atol)
//  [python]
def _test_gelu_op(self, X, fast_gelu, device):
        def _gelu_ref(_X):
            return (_X * norm.cdf(_X).astype(np.float32), )
        expected_output, = _gelu_ref(X)
        actual_output = torch.ops._caffe2.Gelu(torch.tensor(X), fast_gelu)

        rtol = 1e-3 if fast_gelu else 1e-4
        atol = 1e-5
        torch.testing.assert_close(
            expected_output, actual_output.cpu(), rtol=rtol, atol=atol)
//  [python]
def test_bbox_transform(
        self,
        roi_counts,
        num_classes,
        rotated,
        angle_bound_on,
        clip_angle_thresh,
        gc,
        dc,
    ):
        """
        Test with rois for multiple images in a batch
        """
        rois, deltas, im_info = create_bbox_transform_inputs(
            roi_counts, num_classes, rotated
        )

        def bbox_transform_ref():
            ref_op = core.CreateOperator(
                "BBoxTransform",
                ["rois", "deltas", "im_info"],
                ["box_out"],
                apply_scale=False,
                rotated=rotated,
                angle_bound_on=angle_bound_on,
                clip_angle_thresh=clip_angle_thresh,
            )
            workspace.FeedBlob("rois", rois)
            workspace.FeedBlob("deltas", deltas)
            workspace.FeedBlob("im_info", im_info)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("box_out")

        box_out = torch.tensor(bbox_transform_ref())
        a, b = torch.ops._caffe2.BBoxTransform(
            torch.tensor(rois),
            torch.tensor(deltas),
            torch.tensor(im_info),
            [1.0, 1.0, 1.0, 1.0],
            False,
            rotated,
            angle_bound_on,
            -90,
            90,
            clip_angle_thresh,
            legacy_plus_one=True,
        )

        torch.testing.assert_allclose(box_out, a)
//  [python]
def test_bbox_transform(
        self,
        roi_counts,
        num_classes,
        rotated,
        angle_bound_on,
        clip_angle_thresh,
        gc,
        dc,
    ):
        """
        Test with rois for multiple images in a batch
        """
        rois, deltas, im_info = create_bbox_transform_inputs(
            roi_counts, num_classes, rotated
        )

        def bbox_transform_ref():
            ref_op = core.CreateOperator(
                "BBoxTransform",
                ["rois", "deltas", "im_info"],
                ["box_out"],
                apply_scale=False,
                rotated=rotated,
                angle_bound_on=angle_bound_on,
                clip_angle_thresh=clip_angle_thresh,
            )
            workspace.FeedBlob("rois", rois)
            workspace.FeedBlob("deltas", deltas)
            workspace.FeedBlob("im_info", im_info)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("box_out")

        box_out = torch.tensor(bbox_transform_ref())
        a, b = torch.ops._caffe2.BBoxTransform(
            torch.tensor(rois),
            torch.tensor(deltas),
            torch.tensor(im_info),
            [1.0, 1.0, 1.0, 1.0],
            False,
            rotated,
            angle_bound_on,
            -90,
            90,
            clip_angle_thresh,
            legacy_plus_one=True,
        )

        torch.testing.assert_close(box_out, a)
//  [python]
def _test_roi_align(self, N, C, H, W, device):
        def rand_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    0.5 * np.random.rand() * W,
                    0.5 * np.random.rand() * H,
                    (0.5 + 0.5 * np.random.rand()) * W,
                    (0.5 + 0.5 * np.random.rand()) * H,
                ]
            ).astype(np.float32)

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

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlign",
                ["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.RoIAlign(
            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(self, N, C, H, W, device):
        def rand_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    0.5 * np.random.rand() * W,
                    0.5 * np.random.rand() * H,
                    (0.5 + 0.5 * np.random.rand()) * W,
                    (0.5 + 0.5 * np.random.rand()) * H,
                ]
            ).astype(np.float32)

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

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlign",
                ["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.RoIAlign(
            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_lengths_op(self, inputs, ref_op_name, torch_op, device):
        data, lengths = inputs

        def _lengths_ref(X, Y):
            ref_op = core.CreateOperator(ref_op_name, ["X", "Y"], "out")
            workspace.FeedBlob("X", X)
            workspace.FeedBlob("Y", Y)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("out")

        expected_output = _lengths_ref(data, lengths)
        actual_output = torch_op(
            torch.tensor(data), torch.tensor(lengths, dtype=torch.int32))

        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def _test_lengths_op(self, inputs, ref_op_name, torch_op, device):
        data, lengths = inputs

        def _lengths_ref(X, Y):
            ref_op = core.CreateOperator(ref_op_name, ["X", "Y"], "out")
            workspace.FeedBlob("X", X)
            workspace.FeedBlob("Y", Y)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("out")

        expected_output = _lengths_ref(data, lengths)
        actual_output = torch_op(
            torch.tensor(data), torch.tensor(lengths, dtype=torch.int32))

        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_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_alias_with_name_is_in_place(self):
        device = "cuda" if workspace.has_cuda_support else "cpu"
        x = torch.Tensor([3, 42]).to(device)
        y = torch.ops._caffe2.AliasWithName(x, "new_name")
        x[1] = 6
        torch.testing.assert_allclose(x, torch.Tensor([3, 6]).to(device))
        # y should also change because y is alias of x
        torch.testing.assert_allclose(y, torch.Tensor([3, 6]).to(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)
        y = torch.ops._caffe2.AliasWithName(x, "new_name")
        x[1] = 6
        torch.testing.assert_close(x, torch.Tensor([3, 6]).to(device))
        # y should also change because y is alias of x
        torch.testing.assert_close(y, torch.Tensor([3, 6]).to(device))
//  [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_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_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_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_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_percentile(self):
        original_values = np.array([[3., 5., 3], [5., 1., 6.]]).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).int()
        )
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_percentile(self):
        original_values = np.array([[3., 5., 3], [5., 1., 6.]]).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).int()
        )
        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_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).int(), 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).int(), torch.Tensor(boundaries)
        )
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_olmo_model(model_path: str):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from hf_olmo import OLMoForCausalLM, OLMoTokenizerFast  # noqa: F401

    model = OLMo.from_checkpoint(model_path)
    hf_model = AutoModelForCausalLM.from_pretrained(model_path)

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    input = tokenizer.encode("My name is OLMo!")
    input_tensor = torch.tensor(input).unsqueeze(0)

    output = model(input_tensor)
    hf_output = hf_model(input_tensor)

    torch.testing.assert_allclose(output.logits, hf_output.logits)
//  [python]
def test_olmo_model(model_path: str):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from hf_olmo import OLMoForCausalLM, OLMoTokenizerFast  # noqa: F401

    model = OLMo.from_checkpoint(model_path)
    hf_model = AutoModelForCausalLM.from_pretrained(model_path)

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    input = tokenizer.encode("My name is OLMo!")
    input_tensor = torch.tensor(input).unsqueeze(0)

    output = model(input_tensor)
    hf_output = hf_model(input_tensor)

    torch.testing.assert_close(output.logits, hf_output.logits)
//  [python]
def test_save_pretrained(model_path: str):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from hf_olmo import OLMoForCausalLM, OLMoTokenizerFast  # noqa: F401

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    input = tokenizer.encode("My name is OLMo!")
    input_tensor = torch.tensor(input).unsqueeze(0)

    hf_model = AutoModelForCausalLM.from_pretrained(model_path)
    hf_output = hf_model(input_tensor)

    with tempfile.TemporaryDirectory() as tmp_dir:
        hf_model.save_pretrained(tmp_dir)

        saved_hf_model = AutoModelForCausalLM.from_pretrained(tmp_dir)
        saved_hf_output = saved_hf_model(input_tensor)

        torch.testing.assert_allclose(saved_hf_output.logits, hf_output.logits)
//  [python]
def test_save_pretrained(model_path: str):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from hf_olmo import OLMoForCausalLM, OLMoTokenizerFast  # noqa: F401

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    input = tokenizer.encode("My name is OLMo!")
    input_tensor = torch.tensor(input).unsqueeze(0)

    hf_model = AutoModelForCausalLM.from_pretrained(model_path)
    hf_output = hf_model(input_tensor)

    with tempfile.TemporaryDirectory() as tmp_dir:
        hf_model.save_pretrained(tmp_dir)

        saved_hf_model = AutoModelForCausalLM.from_pretrained(tmp_dir)
        saved_hf_output = saved_hf_model(input_tensor)

        torch.testing.assert_close(saved_hf_output.logits, hf_output.logits)
//  [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_set_state_torch(self):
        venv = VectorizedPendulumEnv()
        state, action = self.state_action
        action = np.round(action)

        venv.state = torch.tensor(state)
        vobs, vreward, vdone, _ = venv.step(torch.tensor(action))
        state = venv.state

        venv.set_state(vobs[2:])
        torch.testing.assert_allclose(venv.state, state[2:])
//  [python]
def test_set_state_torch(self):
        venv = VectorizedPendulumEnv()
        state, action = self.state_action
        action = np.round(action)

        venv.state = torch.tensor(state)
        vobs, vreward, vdone, _ = venv.step(torch.tensor(action))
        state = venv.state

        venv.set_state(vobs[2:])
        torch.testing.assert_close(venv.state, state[2:])
//  [python]
def test_pair_batch(self):
        d1 = Data(x=torch.tensor([1]), pos=torch.tensor([1]))
        d2 = Data(x=torch.tensor([2]), pos=torch.tensor([4]))
        d3 = Data(x=torch.tensor([3]), pos=torch.tensor([9]))
        d4 = Data(x=torch.tensor([4]), pos=torch.tensor([16]))
        p1 = Pair.make_pair(d1, d2)
        p2 = Pair.make_pair(d3, d4)
        batch = PairBatch.from_data_list([p1, p2])
        tt.assert_allclose(batch.x, torch.tensor([1, 3]))
        tt.assert_allclose(batch.pos, torch.tensor([1, 9]))
        tt.assert_allclose(batch.batch, torch.tensor([0, 1]))
        tt.assert_allclose(batch.x_target, torch.tensor([2, 4]))
        tt.assert_allclose(batch.pos_target, torch.tensor([4, 16]))
        tt.assert_allclose(batch.batch_target, torch.tensor([0, 1]))
//  [python]
def test_pair_batch(self):
        d1 = Data(x=torch.tensor([1]), pos=torch.tensor([1]))
        d2 = Data(x=torch.tensor([2]), pos=torch.tensor([4]))
        d3 = Data(x=torch.tensor([3]), pos=torch.tensor([9]))
        d4 = Data(x=torch.tensor([4]), pos=torch.tensor([16]))
        p1 = Pair.make_pair(d1, d2)
        p2 = Pair.make_pair(d3, d4)
        batch = PairBatch.from_data_list([p1, p2])
        tt.assert_close(batch.x, torch.tensor([1, 3]))
        tt.assert_close(batch.pos, torch.tensor([1, 9]))
        tt.assert_close(batch.batch, torch.tensor([0, 1]))
        tt.assert_close(batch.x_target, torch.tensor([2, 4]))
        tt.assert_close(batch.pos_target, torch.tensor([4, 16]))
        tt.assert_close(batch.batch_target, torch.tensor([0, 1]))
//  [python]
def test_roll(shifts_dims_pair):
    opset = 11
    from mmcv.onnx.symbolic import register_extra_symbolics
    register_extra_symbolics(opset)

    input = torch.arange(0, 4 * 5 * 6, dtype=torch.float32).view(4, 5, 6)

    shifts, dims = shifts_dims_pair
    func = partial(torch.roll, shifts=shifts, dims=dims)
    wrapped_model = WrapFunction(func).eval()

    with torch.no_grad():
        torch.onnx.export(
            wrapped_model,
            input,
            onnx_file,
            export_params=True,
            keep_initializers_as_inputs=True,
            input_names=['input'],
            output_names=['output'],
            opset_version=opset)

    onnx_model = onnx.load(onnx_file)
    input_all = [node.name for node in onnx_model.graph.input]
    input_initializer = [node.name for node in onnx_model.graph.initializer]
    net_feed_input = list(set(input_all) - set(input_initializer))
    assert (len(net_feed_input) == 1)

    sess = rt.InferenceSession(onnx_file)
    ort_output = sess.run(None, {'input': input.detach().numpy()})[0]

    with torch.no_grad():
        pytorch_output = wrapped_model(input.clone())

    torch.testing.assert_allclose(ort_output, pytorch_output)
//  [python]
def test_roll(shifts_dims_pair):
    opset = 11
    from mmcv.onnx.symbolic import register_extra_symbolics
    register_extra_symbolics(opset)

    input = torch.arange(0, 4 * 5 * 6, dtype=torch.float32).view(4, 5, 6)

    shifts, dims = shifts_dims_pair
    func = partial(torch.roll, shifts=shifts, dims=dims)
    wrapped_model = WrapFunction(func).eval()

    with torch.no_grad():
        torch.onnx.export(
            wrapped_model,
            input,
            onnx_file,
            export_params=True,
            keep_initializers_as_inputs=True,
            input_names=['input'],
            output_names=['output'],
            opset_version=opset)

    onnx_model = onnx.load(onnx_file)
    input_all = [node.name for node in onnx_model.graph.input]
    input_initializer = [node.name for node in onnx_model.graph.initializer]
    net_feed_input = list(set(input_all) - set(input_initializer))
    assert (len(net_feed_input) == 1)

    sess = rt.InferenceSession(onnx_file)
    ort_output = sess.run(None, {'input': input.detach().numpy()})[0]

    with torch.no_grad():
        pytorch_output = wrapped_model(input.clone())

    torch.testing.assert_close(ort_output, pytorch_output)
//  [python]
def _test_gelu_op(self, X, fast_gelu, device):
        def _gelu_ref(_X):
            return (_X * norm.cdf(_X).astype(np.float32),)

        (expected_output,) = _gelu_ref(X)
        actual_output = torch.ops._caffe2.Gelu(torch.tensor(X), fast_gelu)

        rtol = 1e-3 if fast_gelu else 1e-4
        atol = 1e-5
        torch.testing.assert_allclose(
            expected_output, actual_output.cpu(), rtol=rtol, atol=atol
        )
//  [python]
def _test_gelu_op(self, X, fast_gelu, device):
        def _gelu_ref(_X):
            return (_X * norm.cdf(_X).astype(np.float32),)

        (expected_output,) = _gelu_ref(X)
        actual_output = torch.ops._caffe2.Gelu(torch.tensor(X), fast_gelu)

        rtol = 1e-3 if fast_gelu else 1e-4
        atol = 1e-5
        torch.testing.assert_close(
            expected_output, actual_output.cpu(), rtol=rtol, atol=atol
        )
//  [python]
def _test_roi_align(self, N, C, H, W, device):
        def rand_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    0.5 * np.random.rand() * W,
                    0.5 * np.random.rand() * H,
                    (0.5 + 0.5 * np.random.rand()) * W,
                    (0.5 + 0.5 * np.random.rand()) * H,
                ]
            ).astype(np.float32)

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

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlign",
                ["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.RoIAlign(
            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(self, N, C, H, W, device):
        def rand_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    0.5 * np.random.rand() * W,
                    0.5 * np.random.rand() * H,
                    (0.5 + 0.5 * np.random.rand()) * W,
                    (0.5 + 0.5 * np.random.rand()) * H,
                ]
            ).astype(np.float32)

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

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlign",
                ["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.RoIAlign(
            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_lengths_op(self, inputs, ref_op_name, torch_op, device):
        data, lengths = inputs

        def _lengths_ref(X, Y):
            ref_op = core.CreateOperator(ref_op_name, ["X", "Y"], "out")
            workspace.FeedBlob("X", X)
            workspace.FeedBlob("Y", Y)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("out")

        expected_output = _lengths_ref(data, lengths)
        actual_output = torch_op(
            torch.tensor(data), torch.tensor(lengths, dtype=torch.int32)
        )

        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def _test_lengths_op(self, inputs, ref_op_name, torch_op, device):
        data, lengths = inputs

        def _lengths_ref(X, Y):
            ref_op = core.CreateOperator(ref_op_name, ["X", "Y"], "out")
            workspace.FeedBlob("X", X)
            workspace.FeedBlob("Y", Y)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("out")

        expected_output = _lengths_ref(data, lengths)
        actual_output = torch_op(
            torch.tensor(data), torch.tensor(lengths, dtype=torch.int32)
        )

        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [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_allclose(batch.x, torch.tensor([1, 3]))
        tt.assert_allclose(batch.x_target, torch.tensor([2, 4]))
        tt.assert_allclose(batch.batch, torch.tensor([0, 1]))
        tt.assert_allclose(batch.batch_target, torch.tensor([0, 1]))

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

        ms_batches = batch.multiscale_target
        tt.assert_allclose(ms_batches[0].batch, torch.tensor([0, 1]))
        tt.assert_allclose(ms_batches[1].batch, torch.tensor([0, 1]))
        tt.assert_allclose(ms_batches[1].x, torch.tensor([8, 16]))