jiahuimbzuai/CoQuIR huggingface dataset:
//  [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,
        )

        torch.testing.assert_close(box_out, a)
//  [python]
def _test_torch():
    a = torch_compat.asarray([1., 2., 3.])
    b = torch_compat.arange(3, dtype=torch_compat.float64)
    assert a.dtype == torch_compat.float32 == torch.float32
    assert b.dtype == torch_compat.float64 == torch.float64

    # torch.expand_dims does not exist. Update this to use something else if it is added
    res = torch_compat.expand_dims(a, axis=0)
    assert res.dtype == torch_compat.float32 == torch.float32
    assert res.shape == (1, 3)
    assert isinstance(res.shape, torch.Size)
    assert isinstance(a, torch.Tensor)
    assert isinstance(b, torch.Tensor)
    assert isinstance(res, torch.Tensor)

    torch.testing.assert_allclose(res, [[1., 2., 3.]])
//  [python]
def _test_torch():
    a = torch_compat.asarray([1., 2., 3.])
    b = torch_compat.arange(3, dtype=torch_compat.float64)
    assert a.dtype == torch_compat.float32 == torch.float32
    assert b.dtype == torch_compat.float64 == torch.float64

    # torch.expand_dims does not exist. Update this to use something else if it is added
    res = torch_compat.expand_dims(a, axis=0)
    assert res.dtype == torch_compat.float32 == torch.float32
    assert res.shape == (1, 3)
    assert isinstance(res.shape, torch.Size)
    assert isinstance(a, torch.Tensor)
    assert isinstance(b, torch.Tensor)
    assert isinstance(res, torch.Tensor)

    torch.testing.assert_close(res, [[1., 2., 3.]])
//  [python]
def test_constants2():
    constants = [
        sympy.pi,
        sympy.E,
        sympy.GoldenRatio,
        sympy.TribonacciConstant,
        sympy.EulerGamma,
        sympy.Catalan,
    ]
    mod = sympytorch.SymPyModule(expressions=constants)
    mod.to(torch.float64)
    assert mod.sympy() == constants, "mod: {}, y:{}".format(mod.sympy(), constants)
    assert len([p.item() for p in mod.parameters()]) == 0
    torch.testing.assert_allclose(mod(), torch.tensor([float(c) for c in constants]))
//  [python]
def test_constants2():
    constants = [
        sympy.pi,
        sympy.E,
        sympy.GoldenRatio,
        sympy.TribonacciConstant,
        sympy.EulerGamma,
        sympy.Catalan,
    ]
    mod = sympytorch.SymPyModule(expressions=constants)
    mod.to(torch.float64)
    assert mod.sympy() == constants, "mod: {}, y:{}".format(mod.sympy(), constants)
    assert len([p.item() for p in mod.parameters()]) == 0
    torch.testing.assert_close(mod(), torch.tensor([float(c) for c in constants]))
//  [python]
def test_tabular_interaction(agent, policy):
    LEARNING_RATE = 0.1
    environment = EasyGridWorld()

    critic = TabularQFunction(
        num_states=environment.num_states, num_actions=environment.num_actions
    )
    policy = policy(critic, 0.1)
    optimizer = torch.optim.Adam(critic.parameters(), lr=LEARNING_RATE)
    criterion = torch.nn.MSELoss
    memory = ExperienceReplay(max_len=MEMORY_MAX_SIZE)

    agent = agent(
        critic=critic,
        policy=policy,
        criterion=criterion,
        optimizer=optimizer,
        memory=memory,
        batch_size=BATCH_SIZE,
        target_update_frequency=TARGET_UPDATE_FREQUENCY,
        gamma=GAMMA,
    )

    train_agent(
        agent,
        environment,
        num_episodes=NUM_EPISODES,
        max_steps=MAX_STEPS,
        plot_flag=False,
    )
    evaluate_agent(agent, environment, 1, MAX_STEPS, render=False)
    agent.logger.delete_directory()  # Cleanup directory.

    torch.testing.assert_allclose(
        critic.table.shape,
        torch.Size([environment.num_actions, environment.num_states, 1]),
    )
//  [python]
def test_tabular_interaction(agent, policy):
    LEARNING_RATE = 0.1
    environment = EasyGridWorld()

    critic = TabularQFunction(
        num_states=environment.num_states, num_actions=environment.num_actions
    )
    policy = policy(critic, 0.1)
    optimizer = torch.optim.Adam(critic.parameters(), lr=LEARNING_RATE)
    criterion = torch.nn.MSELoss
    memory = ExperienceReplay(max_len=MEMORY_MAX_SIZE)

    agent = agent(
        critic=critic,
        policy=policy,
        criterion=criterion,
        optimizer=optimizer,
        memory=memory,
        batch_size=BATCH_SIZE,
        target_update_frequency=TARGET_UPDATE_FREQUENCY,
        gamma=GAMMA,
    )

    train_agent(
        agent,
        environment,
        num_episodes=NUM_EPISODES,
        max_steps=MAX_STEPS,
        plot_flag=False,
    )
    evaluate_agent(agent, environment, 1, MAX_STEPS, render=False)
    agent.logger.delete_directory()  # Cleanup directory.

    torch.testing.assert_close(
        critic.table.shape,
        torch.Size([environment.num_actions, environment.num_states, 1]),
    )
//  [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,
        )
        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,
        )
        torch.testing.assert_close(roi_feature_ref, roi_feature.cpu())
//  [python]
def test_project(self) -> None:
        device = torch.device("cpu")
        dtype = torch.float32
        camera = PerspectiveCameras(device=device)
        mask = torch.ones(2, 1, 4, 6, device=device, dtype=dtype)
        src = GridImage(
            data=torch.ones(2, 3, 4, 6, device=device, dtype=dtype),
            camera=camera,
            time=torch.rand(2, device=device),
            mask=mask,
        )
        depth = GridImage(
            data=torch.ones(2, 1, 4, 6, device=device, dtype=dtype),
            camera=camera,
            time=torch.rand(2, device=device),
            mask=mask,
        )
        vel = GridImage(
            data=torch.zeros(2, 3, 4, 6, device=device, dtype=dtype),
            camera=camera,
            time=torch.rand(2, device=device),
            mask=mask,
        )
        # compiled_project = torch.compile(project, fullgraph=True, backend="eager")
        proj = project(dst=src, src=src, depth=depth, vel=vel)
        torch.testing.assert_allclose(proj._data, src._data)
        torch.testing.assert_allclose(proj.mask._data, mask)
//  [python]
def test_project(self) -> None:
        device = torch.device("cpu")
        dtype = torch.float32
        camera = PerspectiveCameras(device=device)
        mask = torch.ones(2, 1, 4, 6, device=device, dtype=dtype)
        src = GridImage(
            data=torch.ones(2, 3, 4, 6, device=device, dtype=dtype),
            camera=camera,
            time=torch.rand(2, device=device),
            mask=mask,
        )
        depth = GridImage(
            data=torch.ones(2, 1, 4, 6, device=device, dtype=dtype),
            camera=camera,
            time=torch.rand(2, device=device),
            mask=mask,
        )
        vel = GridImage(
            data=torch.zeros(2, 3, 4, 6, device=device, dtype=dtype),
            camera=camera,
            time=torch.rand(2, device=device),
            mask=mask,
        )
        # compiled_project = torch.compile(project, fullgraph=True, backend="eager")
        proj = project(dst=src, src=src, depth=depth, vel=vel)
        torch.testing.assert_close(proj._data, src._data)
        torch.testing.assert_close(proj.mask._data, mask)
//  [python]
x1 = torch.tensor(__example_input_array__, requires_grad=True)
y1 = m(x1)
y1.backward()
print(x1.grad)
torch.testing.assert_allclose(x1.grad, torch.tensor(__example_output_array__))
//  [python]
x1 = torch.tensor(__example_input_array__, requires_grad=True)
y1 = m(x1)
y1.backward()
print(x1.grad)
torch.testing.assert_close(x1.grad, torch.tensor(__example_output_array__))
//  [python]
def test_k_winners2d_global_grad(self):
        """
        Test gradient
        """
        x = self.x2.clone().requires_grad_(True)
        n, c, h, w = x.shape

        grad = self.gradient2

        expected = torch.zeros_like(x)
        expected[0, 0, 1, 0] = grad[0, 0, 1, 0]
        expected[0, 0, 1, 1] = grad[0, 0, 1, 1]
        expected[0, 1, 0, 1] = grad[0, 1, 0, 1]
        expected[0, 2, 1, 0] = grad[0, 2, 1, 0]

        expected[1, 0, 0, 0] = grad[1, 0, 0, 0]
        expected[1, 1, 0, 0] = grad[1, 1, 0, 0]
        expected[1, 1, 0, 1] = grad[1, 1, 0, 1]
        expected[1, 2, 1, 1] = grad[1, 2, 1, 1]

        for break_ties in [True, False]:
            with self.subTest(break_ties=break_ties):
                kw = KWinners2d(
                    percent_on=(1 / 3),
                    channels=c,
                    k_inference_factor=1.0,
                    boost_strength=0.0,
                    duty_cycle_period=1000,
                    local=False,
                    break_ties=break_ties,
                )
                kw.train(mode=True)
                y = kw(x)
                y.backward(grad)
                torch.testing.assert_allclose(x.grad, expected)
                x.grad.zero_()
//  [python]
def test_k_winners2d_global_grad(self):
        """
        Test gradient
        """
        x = self.x2.clone().requires_grad_(True)
        n, c, h, w = x.shape

        grad = self.gradient2

        expected = torch.zeros_like(x)
        expected[0, 0, 1, 0] = grad[0, 0, 1, 0]
        expected[0, 0, 1, 1] = grad[0, 0, 1, 1]
        expected[0, 1, 0, 1] = grad[0, 1, 0, 1]
        expected[0, 2, 1, 0] = grad[0, 2, 1, 0]

        expected[1, 0, 0, 0] = grad[1, 0, 0, 0]
        expected[1, 1, 0, 0] = grad[1, 1, 0, 0]
        expected[1, 1, 0, 1] = grad[1, 1, 0, 1]
        expected[1, 2, 1, 1] = grad[1, 2, 1, 1]

        for break_ties in [True, False]:
            with self.subTest(break_ties=break_ties):
                kw = KWinners2d(
                    percent_on=(1 / 3),
                    channels=c,
                    k_inference_factor=1.0,
                    boost_strength=0.0,
                    duty_cycle_period=1000,
                    local=False,
                    break_ties=break_ties,
                )
                kw.train(mode=True)
                y = kw(x)
                y.backward(grad)
                torch.testing.assert_close(x.grad, expected)
                x.grad.zero_()
//  [python]
def test_check_input_permutation(af_type: acquisition_functions.AcquisitionFunction):
    if af_type == acquisition_functions.AcquisitionFunction.random:
        return

    batch_size = 12

    test_data = torch.rand((batch_size, 10))

    mixture_a = test_data[::2, :]
    mixture_b = test_data[1::2, :]
    mixture_c = test_data

    class Forwarder(torch.nn.Module):
        def forward(self, batch):
            return batch

    forwarder = Forwarder()
    estimator = af_type.create(forwarder, k=1)
    estimator.eval()

    output_a = estimator(mixture_a)
    output_b = estimator(mixture_b)
    output_c = estimator(mixture_c)

    torch.testing.assert_allclose(
        torch.cat([output_a, output_b], dim=0), torch.cat([output_c[::2], output_c[1::2]], dim=0)
    )
//  [python]
def test_check_input_permutation(af_type: acquisition_functions.AcquisitionFunction):
    if af_type == acquisition_functions.AcquisitionFunction.random:
        return

    batch_size = 12

    test_data = torch.rand((batch_size, 10))

    mixture_a = test_data[::2, :]
    mixture_b = test_data[1::2, :]
    mixture_c = test_data

    class Forwarder(torch.nn.Module):
        def forward(self, batch):
            return batch

    forwarder = Forwarder()
    estimator = af_type.create(forwarder, k=1)
    estimator.eval()

    output_a = estimator(mixture_a)
    output_b = estimator(mixture_b)
    output_c = estimator(mixture_c)

    torch.testing.assert_close(
        torch.cat([output_a, output_b], dim=0), torch.cat([output_c[::2], output_c[1::2]], dim=0)
    )
//  [python]
def run_test(self, mod, inputs, expected_ops, unexpected_ops, interpreter, rtol, atol, precision=LowerPrecision.FP32):
        with torch.no_grad():
            cuda_inputs = []
            for i in inputs:
                cuda_inputs.append(i.cuda())

            mod.eval()
            if len(expected_ops):
                self.assert_has_op(mod, expected_ops)
            if unexpected_ops:
                self.assert_unexpected_op(mod, unexpected_ops)

            interpreter_result = interpreter.run(lower_precision=precision)
            trt_mod = TRTModule(
                interpreter_result.engine,
                interpreter_result.input_names,
                interpreter_result.output_names,
            )

            ref_outputs = mod(*inputs)
            outputs = trt_mod(*cuda_inputs)

            if isinstance(outputs, torch.Tensor):
                ref_outputs = [ref_outputs]
                outputs = [outputs]
            for out, ref in zip(outputs, ref_outputs):
                if not isinstance(ref, torch.Tensor):
                    ref = torch.tensor([ref])
                ref = ref.cpu()  # to_dtype test has cases with gpu output
                torch.testing.assert_allclose(out.cpu(), ref, rtol=rtol, atol=atol)
//  [python]
def run_test(self, mod, inputs, expected_ops, unexpected_ops, interpreter, rtol, atol, precision=LowerPrecision.FP32):
        with torch.no_grad():
            cuda_inputs = []
            for i in inputs:
                cuda_inputs.append(i.cuda())

            mod.eval()
            if len(expected_ops):
                self.assert_has_op(mod, expected_ops)
            if unexpected_ops:
                self.assert_unexpected_op(mod, unexpected_ops)

            interpreter_result = interpreter.run(lower_precision=precision)
            trt_mod = TRTModule(
                interpreter_result.engine,
                interpreter_result.input_names,
                interpreter_result.output_names,
            )

            ref_outputs = mod(*inputs)
            outputs = trt_mod(*cuda_inputs)

            if isinstance(outputs, torch.Tensor):
                ref_outputs = [ref_outputs]
                outputs = [outputs]
            for out, ref in zip(outputs, ref_outputs):
                if not isinstance(ref, torch.Tensor):
                    ref = torch.tensor([ref])
                ref = ref.cpu()  # to_dtype test has cases with gpu output
                torch.testing.assert_close(out.cpu(), ref, rtol=rtol, atol=atol)
//  [python]
def test_transform_matrix_conversion_functionality(self) -> None:
        """
        Test the numerical accuracy of the transform matrix conversion utilities.
        """
        initial_state = torch.tensor([5, 6, math.pi / 2], dtype=torch.float32)

        expected_xform_matrix = torch.tensor([[0, -1, 5], [1, 0, 6], [0, 0, 1]], dtype=torch.float32)

        xform_matrix = state_se2_tensor_to_transform_matrix(initial_state, precision=torch.float32)

        torch.testing.assert_allclose(expected_xform_matrix, xform_matrix)

        reverted = transform_matrix_to_state_se2_tensor(xform_matrix, precision=torch.float32)

        torch.testing.assert_allclose(initial_state, reverted)
//  [python]
def test_transform_matrix_conversion_functionality(self) -> None:
        """
        Test the numerical accuracy of the transform matrix conversion utilities.
        """
        initial_state = torch.tensor([5, 6, math.pi / 2], dtype=torch.float32)

        expected_xform_matrix = torch.tensor([[0, -1, 5], [1, 0, 6], [0, 0, 1]], dtype=torch.float32)

        xform_matrix = state_se2_tensor_to_transform_matrix(initial_state, precision=torch.float32)

        torch.testing.assert_close(expected_xform_matrix, xform_matrix)

        reverted = transform_matrix_to_state_se2_tensor(xform_matrix, precision=torch.float32)

        torch.testing.assert_close(initial_state, reverted)
//  [python]
def test_transform_matrix_scriptability(self) -> None:
        """
        Tests that the transform matrix conversion utilities script properly.
        """

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

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                xform = state_se2_tensor_to_transform_matrix(x)
                result = transform_matrix_to_state_se2_tensor(xform)
                return result

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

        test_input = torch.tensor([1, 2, 3], dtype=torch.float32)

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

        torch.testing.assert_allclose(py_result, script_result)
//  [python]
def test_transform_matrix_scriptability(self) -> None:
        """
        Tests that the transform matrix conversion utilities script properly.
        """

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

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                xform = state_se2_tensor_to_transform_matrix(x)
                result = transform_matrix_to_state_se2_tensor(xform)
                return result

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

        test_input = torch.tensor([1, 2, 3], dtype=torch.float32)

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

        torch.testing.assert_close(py_result, script_result)
//  [python]
def test_get_group_qparams_symmetric(self):
        """
        Test that `get_group_qparams_symmetric` produces the exact same scales as
        `PerChannelMinMaxObserver._calculate_qparams`.
        """
        n_bit = 4
        qmin = -(2 ** (n_bit - 1))
        qmax = 2 ** (n_bit - 1) - 1
        eps = torch.finfo(torch.float32).eps
        groupsize = 256
        torch.manual_seed(self.SEED)
        weight = torch.randn(100, 256).to(torch.float16)

        # calculate observer scales
        obs = torch.ao.quantization.PerChannelMinMaxObserver(
            ch_axis=0,
            qscheme=torch.per_channel_symmetric,
            quant_min=qmin,
            quant_max=qmax,
            # This is needed to ensure `min_val` and `max_val` are fp16,
            # otherwise they default to fp32 and the qparams will be slightly off
            factory_kwargs={"dtype": torch.float16}
        )
        obs(weight)
        (scale_obs, _) = obs.calculate_qparams()
        scale_obs = scale_obs.reshape(weight.shape[0], -1)

        # assert that scales are identical
        (scale_ao, _) = get_group_qparams_symmetric(weight, n_bit, groupsize)
        torch.testing.assert_allclose(scale_obs, scale_ao, rtol=0, atol=0)
//  [python]
def test_get_group_qparams_symmetric(self):
        """
        Test that `get_group_qparams_symmetric` produces the exact same scales as
        `PerChannelMinMaxObserver._calculate_qparams`.
        """
        n_bit = 4
        qmin = -(2 ** (n_bit - 1))
        qmax = 2 ** (n_bit - 1) - 1
        eps = torch.finfo(torch.float32).eps
        groupsize = 256
        torch.manual_seed(self.SEED)
        weight = torch.randn(100, 256).to(torch.float16)

        # calculate observer scales
        obs = torch.ao.quantization.PerChannelMinMaxObserver(
            ch_axis=0,
            qscheme=torch.per_channel_symmetric,
            quant_min=qmin,
            quant_max=qmax,
            # This is needed to ensure `min_val` and `max_val` are fp16,
            # otherwise they default to fp32 and the qparams will be slightly off
            factory_kwargs={"dtype": torch.float16}
        )
        obs(weight)
        (scale_obs, _) = obs.calculate_qparams()
        scale_obs = scale_obs.reshape(weight.shape[0], -1)

        # assert that scales are identical
        (scale_ao, _) = get_group_qparams_symmetric(weight, n_bit, groupsize)
        torch.testing.assert_close(scale_obs, scale_ao, rtol=0, atol=0)
//  [python]
def test_transform_matrix_batch_conversion_functionality(self) -> None:
        """
        Test the numerical accuracy of the transform matrix conversion utilities.
        """
        initial_state = torch.tensor([[5, 6, math.pi / 2]], dtype=torch.float32)

        expected_xform_matrix = torch.tensor([[[0, -1, 5], [1, 0, 6], [0, 0, 1]]], dtype=torch.float32)

        xform_matrix = state_se2_tensor_to_transform_matrix_batch(initial_state, precision=torch.float32)

        torch.testing.assert_allclose(expected_xform_matrix, xform_matrix)

        reverted = transform_matrix_to_state_se2_tensor_batch(xform_matrix)

        torch.testing.assert_allclose(initial_state, reverted)

        with self.assertRaises(ValueError):
            # We only accept Nx3x3, not 3x3 tensors.
            misshaped_tensor = torch.tensor([1, 2, 3], dtype=torch.float32)
            state_se2_tensor_to_transform_matrix_batch(misshaped_tensor)
//  [python]
def test_transform_matrix_batch_conversion_functionality(self) -> None:
        """
        Test the numerical accuracy of the transform matrix conversion utilities.
        """
        initial_state = torch.tensor([[5, 6, math.pi / 2]], dtype=torch.float32)

        expected_xform_matrix = torch.tensor([[[0, -1, 5], [1, 0, 6], [0, 0, 1]]], dtype=torch.float32)

        xform_matrix = state_se2_tensor_to_transform_matrix_batch(initial_state, precision=torch.float32)

        torch.testing.assert_close(expected_xform_matrix, xform_matrix)

        reverted = transform_matrix_to_state_se2_tensor_batch(xform_matrix)

        torch.testing.assert_close(initial_state, reverted)

        with self.assertRaises(ValueError):
            # We only accept Nx3x3, not 3x3 tensors.
            misshaped_tensor = torch.tensor([1, 2, 3], dtype=torch.float32)
            state_se2_tensor_to_transform_matrix_batch(misshaped_tensor)
//  [python]
def test_transform_matrix_batch_scriptability(self) -> None:
        """
        Tests that the transform matrix conversion utilities script properly.
        """

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

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                xform = state_se2_tensor_to_transform_matrix_batch(x)
                result = transform_matrix_to_state_se2_tensor_batch(xform)
                return result

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

        test_input = torch.tensor([[1, 2, 3]], dtype=torch.float32)

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

        torch.testing.assert_allclose(py_result, script_result)
//  [python]
def test_transform_matrix_batch_scriptability(self) -> None:
        """
        Tests that the transform matrix conversion utilities script properly.
        """

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

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                xform = state_se2_tensor_to_transform_matrix_batch(x)
                result = transform_matrix_to_state_se2_tensor_batch(xform)
                return result

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

        test_input = torch.tensor([[1, 2, 3]], dtype=torch.float32)

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

        torch.testing.assert_close(py_result, script_result)
//  [python]
def test_global_state_se2_tensor_to_local_scriptability(self) -> None:
        """
        Tests that global_state_se2_tensor_to_local scripts properly.
        """

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

            def forward(self, states: torch.Tensor, pose: torch.Tensor) -> torch.Tensor:
                result = global_state_se2_tensor_to_local(states, pose)
                return result

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

        test_states = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
        test_pose = torch.tensor([1, 2, 3], dtype=torch.float32)

        py_result = to_script.forward(test_states, test_pose)
        script_result = scripted.forward(test_states, test_pose)

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

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

            def forward(self, states: torch.Tensor, pose: torch.Tensor) -> torch.Tensor:
                result = global_state_se2_tensor_to_local(states, pose)
                return result

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

        test_states = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
        test_pose = torch.tensor([1, 2, 3], dtype=torch.float32)

        py_result = to_script.forward(test_states, test_pose)
        script_result = scripted.forward(test_states, test_pose)

        torch.testing.assert_close(py_result, script_result)
//  [python]
def test_global_state_se2_tensor_to_local_functionality(self) -> None:
        """
        Tests the numerical accuracy of global_state_se2_tensor_to_local.
        """
        """
           o = coordinates to transform, facing direction >
           # = local reference frame
           y_world
            ^              x_local
            |               ^
            |               |
            |   y_local <---#
            |
            |
            |
         o  |  o>
         V  |
            |
            *---------------------> x_world
               ^
        <o     o

        """
        global_states = torch.tensor(
            [[1, 1, 0], [1, -1, math.pi / 2], [-1, -1, math.pi], [-1, 1, -math.pi / 2]], dtype=torch.float32
        )

        local_state = torch.tensor([5, 5, math.pi / 2], dtype=torch.float32)

        expected_transformed_states = torch.tensor(
            [[-4, 4, -math.pi / 2], [-6, 4, 0], [-6, 6, math.pi / 2], [-4, 6, math.pi]], dtype=torch.float32
        )

        actual_transformed_states = global_state_se2_tensor_to_local(
            global_states, local_state, precision=torch.float32
        )

        torch.testing.assert_allclose(expected_transformed_states, actual_transformed_states)
//  [python]
def test_global_state_se2_tensor_to_local_functionality(self) -> None:
        """
        Tests the numerical accuracy of global_state_se2_tensor_to_local.
        """
        """
           o = coordinates to transform, facing direction >
           # = local reference frame
           y_world
            ^              x_local
            |               ^
            |               |
            |   y_local <---#
            |
            |
            |
         o  |  o>
         V  |
            |
            *---------------------> x_world
               ^
        <o     o

        """
        global_states = torch.tensor(
            [[1, 1, 0], [1, -1, math.pi / 2], [-1, -1, math.pi], [-1, 1, -math.pi / 2]], dtype=torch.float32
        )

        local_state = torch.tensor([5, 5, math.pi / 2], dtype=torch.float32)

        expected_transformed_states = torch.tensor(
            [[-4, 4, -math.pi / 2], [-6, 4, 0], [-6, 6, math.pi / 2], [-4, 6, math.pi]], dtype=torch.float32
        )

        actual_transformed_states = global_state_se2_tensor_to_local(
            global_states, local_state, precision=torch.float32
        )

        torch.testing.assert_close(expected_transformed_states, actual_transformed_states)
//  [python]
def test_coordinates_to_local_frame_functionality(self) -> None:
        """
        Tests the numerical accuracy of coordinates_to_local_frame.
        """
        """
           o = coordinates to transform
           # = local reference frame
           y_world
            ^              x_local
            |               ^
            |               |
            |   y_local <---#
            |
            |
            |
         o  |  o
            |
            |
            *---------------------> x_world

         o     o

        """
        coordinates = torch.tensor([[1, 1], [1, -1], [-1, -1], [-1, 1]], dtype=torch.float32)

        local_state = torch.tensor([5, 5, math.pi / 2], dtype=torch.float32)

        expected_coordinates = torch.tensor([[-4, 4], [-6, 4], [-6, 6], [-4, 6]], dtype=torch.float32)

        actual_coordinates = coordinates_to_local_frame(coordinates, local_state, precision=torch.float32)

        torch.testing.assert_allclose(expected_coordinates, actual_coordinates)
//  [python]
def test_coordinates_to_local_frame_functionality(self) -> None:
        """
        Tests the numerical accuracy of coordinates_to_local_frame.
        """
        """
           o = coordinates to transform
           # = local reference frame
           y_world
            ^              x_local
            |               ^
            |               |
            |   y_local <---#
            |
            |
            |
         o  |  o
            |
            |
            *---------------------> x_world

         o     o

        """
        coordinates = torch.tensor([[1, 1], [1, -1], [-1, -1], [-1, 1]], dtype=torch.float32)

        local_state = torch.tensor([5, 5, math.pi / 2], dtype=torch.float32)

        expected_coordinates = torch.tensor([[-4, 4], [-6, 4], [-6, 6], [-4, 6]], dtype=torch.float32)

        actual_coordinates = coordinates_to_local_frame(coordinates, local_state, precision=torch.float32)

        torch.testing.assert_close(expected_coordinates, actual_coordinates)
//  [python]
def test_coordinates_to_local_frame_scriptability(self) -> None:
        """
        Tests that the function coordinates_to_local_frame scripts properly.
        """

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

            def forward(self, states: torch.Tensor, pose: torch.Tensor) -> torch.Tensor:
                result = coordinates_to_local_frame(states, pose, precision=torch.float32)
                return result

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

        test_states = torch.tensor([[1, 2], [4, 5]], dtype=torch.float32)
        test_pose = torch.tensor([1, 2, 3], dtype=torch.float32)

        py_result = to_script.forward(test_states, test_pose)
        script_result = scripted.forward(test_states, test_pose)

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

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

            def forward(self, states: torch.Tensor, pose: torch.Tensor) -> torch.Tensor:
                result = coordinates_to_local_frame(states, pose, precision=torch.float32)
                return result

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

        test_states = torch.tensor([[1, 2], [4, 5]], dtype=torch.float32)
        test_pose = torch.tensor([1, 2, 3], dtype=torch.float32)

        py_result = to_script.forward(test_states, test_pose)
        script_result = scripted.forward(test_states, test_pose)

        torch.testing.assert_close(py_result, script_result)
//  [python]
lowered = lower_to_elementwise_interpreter(mem)
print(lowered.code)
# The lowered module can also be compiled into TorchScript
scripted = torch.jit.script(lowered)
print(scripted.graph)

# Stress test correctness
for _ in range(50):
    x, y = torch.randn(10, 20, 30), torch.randn(10, 20, 30)
    torch.testing.assert_allclose(lowered(x, y), mem(x, y))
    torch.testing.assert_allclose(scripted(x, y), mem(x, y))
//  [python]
lowered = lower_to_elementwise_interpreter(mem)
print(lowered.code)
# The lowered module can also be compiled into TorchScript
scripted = torch.jit.script(lowered)
print(scripted.graph)

# Stress test correctness
for _ in range(50):
    x, y = torch.randn(10, 20, 30), torch.randn(10, 20, 30)
    torch.testing.assert_close(lowered(x, y), mem(x, y))
    torch.testing.assert_close(scripted(x, y), mem(x, y))
//  [python]
def test_vector_set_coordinates_to_local_frame_functionality(self) -> None:
        """
        Test converting vector set map coordinates from global to local ego frame.
        """
        coords = torch.tensor([[[1, 1], [3, 1], [5, 1]]], dtype=torch.float64)
        avails = torch.ones(coords.shape[:2], dtype=torch.bool)
        anchor_state = torch.tensor([0, 0, 0], dtype=torch.float64)

        result_coords = vector_set_coordinates_to_local_frame(coords, avails, anchor_state)

        self.assertIsInstance(result_coords, torch.FloatTensor)
        self.assertEqual(result_coords.shape, coords.shape)
        torch.testing.assert_allclose(coords.float(), result_coords)

        # test unexpected shape
        with self.assertRaises(ValueError):
            vector_set_coordinates_to_local_frame(coords[0], avails[0], anchor_state)

        # test mismatching shape
        with self.assertRaises(ValueError):
            vector_set_coordinates_to_local_frame(coords, avails[0], anchor_state)
//  [python]
def test_vector_set_coordinates_to_local_frame_functionality(self) -> None:
        """
        Test converting vector set map coordinates from global to local ego frame.
        """
        coords = torch.tensor([[[1, 1], [3, 1], [5, 1]]], dtype=torch.float64)
        avails = torch.ones(coords.shape[:2], dtype=torch.bool)
        anchor_state = torch.tensor([0, 0, 0], dtype=torch.float64)

        result_coords = vector_set_coordinates_to_local_frame(coords, avails, anchor_state)

        self.assertIsInstance(result_coords, torch.FloatTensor)
        self.assertEqual(result_coords.shape, coords.shape)
        torch.testing.assert_close(coords.float(), result_coords)

        # test unexpected shape
        with self.assertRaises(ValueError):
            vector_set_coordinates_to_local_frame(coords[0], avails[0], anchor_state)

        # test mismatching shape
        with self.assertRaises(ValueError):
            vector_set_coordinates_to_local_frame(coords, avails[0], anchor_state)
//  [python]
def test_vector_set_coordinates_to_local_frame_scriptability(self) -> None:
        """
        Tests that the function vector_set_coordinates_to_local_frame scripts properly.
        """

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

            def forward(self, coords: torch.Tensor, avails: torch.Tensor, anchor_state: torch.Tensor) -> torch.Tensor:
                result = vector_set_coordinates_to_local_frame(coords, avails, anchor_state)
                return result

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

        test_coords = torch.tensor([[[1, 1], [3, 1], [5, 1]]], dtype=torch.float64)
        test_avails = torch.ones(test_coords.shape[:2], dtype=torch.bool)
        test_anchor_state = torch.tensor([0, 0, 0], dtype=torch.float64)

        py_result = to_script.forward(test_coords, test_avails, test_anchor_state)
        script_result = scripted.forward(test_coords, test_avails, test_anchor_state)

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

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

            def forward(self, coords: torch.Tensor, avails: torch.Tensor, anchor_state: torch.Tensor) -> torch.Tensor:
                result = vector_set_coordinates_to_local_frame(coords, avails, anchor_state)
                return result

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

        test_coords = torch.tensor([[[1, 1], [3, 1], [5, 1]]], dtype=torch.float64)
        test_avails = torch.ones(test_coords.shape[:2], dtype=torch.bool)
        test_anchor_state = torch.tensor([0, 0, 0], dtype=torch.float64)

        py_result = to_script.forward(test_coords, test_avails, test_anchor_state)
        script_result = scripted.forward(test_coords, test_avails, test_anchor_state)

        torch.testing.assert_close(py_result, script_result)
//  [python]
def test_scalePos(self):
        tr = ScalePos(scale=2.0)
        d = Data(pos=torch.tensor([[1, 0, 0], [0, 1, 1]]).float())
        d = tr(d)
        torch.testing.assert_allclose(d.pos, torch.tensor([[2, 0, 0], [0, 2, 2]]).float())
//  [python]
def test_scalePos(self):
        tr = ScalePos(scale=2.0)
        d = Data(pos=torch.tensor([[1, 0, 0], [0, 1, 1]]).float())
        d = tr(d)
        torch.testing.assert_close(d.pos, torch.tensor([[2, 0, 0], [0, 2, 2]]).float())
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 64, 64)
        data_samples = [ClsDataSample().set_gt_label([1, 2, 6])]
        # default
        model = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_allclose(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # k is not -1
        cfg = {**self.DEFAULT_ARGS, 'topk': 2}
        model = MODELS.build(cfg)

        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        assert predictions is data_samples
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 64, 64)
        data_samples = [ClsDataSample().set_gt_label([1, 2, 6])]
        # default
        model = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_close(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # k is not -1
        cfg = {**self.DEFAULT_ARGS, 'topk': 2}
        model = MODELS.build(cfg)

        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        assert predictions is data_samples
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
//  [python]
def test_normalizeFeature(self):
        tr = NormalizeFeature(feature_name="new_feature", standardize=False)
        d = Data(new_feature=torch.tensor([[-3, -2, -1, 0, 1, 2]]).float())
        d = tr(d)
        torch.testing.assert_allclose(d.new_feature, torch.tensor([[0, 0.2, 0.4, 0.6, 0.8, 1]]).float())

        tr = NormalizeFeature(feature_name="new_feature", standardize=True)
        d = Data(new_feature=torch.tensor([[0, 1]]).float())
        d = tr(d)
        torch.testing.assert_allclose(d.new_feature, (d.new_feature - d.new_feature.mean()) / d.new_feature.std())
//  [python]
def test_normalizeFeature(self):
        tr = NormalizeFeature(feature_name="new_feature", standardize=False)
        d = Data(new_feature=torch.tensor([[-3, -2, -1, 0, 1, 2]]).float())
        d = tr(d)
        torch.testing.assert_close(d.new_feature, torch.tensor([[0, 0.2, 0.4, 0.6, 0.8, 1]]).float())

        tr = NormalizeFeature(feature_name="new_feature", standardize=True)
        d = Data(new_feature=torch.tensor([[0, 1]]).float())
        d = tr(d)
        torch.testing.assert_close(d.new_feature, (d.new_feature - d.new_feature.mean()) / d.new_feature.std())
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 64, 64)
        data_samples = [ClsDataSample().set_gt_label(1)]
        model = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        # assert False, type(outs)
        self.assertIsInstance(outs, tuple)
        self.assertEqual(len(outs), 1)
        self.assertIsInstance(outs[0], torch.Tensor)

        # test forward train
        losses = model(inputs, data_samples, mode='loss')
        self.assertGreater(losses['loss'].item(), 0)

        # test forward test
        predictions = model(inputs, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_allclose(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 64, 64)
        data_samples = [ClsDataSample().set_gt_label(1)]
        model = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        # assert False, type(outs)
        self.assertIsInstance(outs, tuple)
        self.assertEqual(len(outs), 1)
        self.assertIsInstance(outs[0], torch.Tensor)

        # test forward train
        losses = model(inputs, data_samples, mode='loss')
        self.assertGreater(losses['loss'].item(), 0)

        # test forward test
        predictions = model(inputs, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_close(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_discrete_distribution(self):
        d = Categorical(torch.tensor([0.1, 0.2, 0.3, 0.4]))

        def _function(a):
            return 2 * a

        torch.testing.assert_allclose(integrate(_function, d), 4.0)
//  [python]
def test_discrete_distribution(self):
        d = Categorical(torch.tensor([0.1, 0.2, 0.3, 0.4]))

        def _function(a):
            return 2 * a

        torch.testing.assert_close(integrate(_function, d), 4.0)
//  [python]
def test_select(self):
        data = Data(pos=torch.tensor([0, 1, 2]), x=torch.tensor([2, 3, 4]), z=torch.tensor([1]))
        tr = Select(torch.tensor([1, 2]))

        data2 = tr(data)
        torch.testing.assert_allclose(data2.pos, torch.tensor([1, 2]))
        torch.testing.assert_allclose(data2.x, torch.tensor([3, 4]))
        torch.testing.assert_allclose(data2.z, torch.tensor([1]))

        tr = Select(torch.tensor([False, True, True]))

        data2 = tr(data)
        torch.testing.assert_allclose(data2.pos, torch.tensor([1, 2]))
        torch.testing.assert_allclose(data2.x, torch.tensor([3, 4]))
        torch.testing.assert_allclose(data2.z, torch.tensor([1]))
//  [python]
def test_select(self):
        data = Data(pos=torch.tensor([0, 1, 2]), x=torch.tensor([2, 3, 4]), z=torch.tensor([1]))
        tr = Select(torch.tensor([1, 2]))

        data2 = tr(data)
        torch.testing.assert_close(data2.pos, torch.tensor([1, 2]))
        torch.testing.assert_close(data2.x, torch.tensor([3, 4]))
        torch.testing.assert_close(data2.z, torch.tensor([1]))

        tr = Select(torch.tensor([False, True, True]))

        data2 = tr(data)
        torch.testing.assert_close(data2.pos, torch.tensor([1, 2]))
        torch.testing.assert_close(data2.x, torch.tensor([3, 4]))
        torch.testing.assert_close(data2.z, torch.tensor([1]))