jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_dfop():
    from df.config import config

    config.use_defaults()
    p = ModelParams()
    f = p.nb_df
    F = f * 2
    o = p.df_order
    d = p.df_lookahead
    t = 100
    spec = torch.randn(1, 1, t, F, 2)
    coefs = torch.randn(1, t, o, f, 2)
    alpha = torch.randn(1, t, 1)
    dfop = DfOp(df_bins=p.nb_df)
    dfop.set_forward("real_loop")
    out1 = dfop(spec, coefs, alpha)
    dfop.set_forward("real_strided")
    out2 = dfop(spec, coefs, alpha)
    dfop.set_forward("real_unfold")
    out3 = dfop(spec, coefs, alpha)
    dfop.set_forward("complex_strided")
    out4 = dfop(spec, coefs, alpha)
    torch.testing.assert_allclose(out1, out2)
    torch.testing.assert_allclose(out1, out3)
    torch.testing.assert_allclose(out1, out4)
    # This forward method requires external padding/lookahead as well as spectrogram buffer
    # handling, i.e. via a ring buffer. Could be used in real time usage.
    dfop.set_forward("real_one_step")
    spec_padded = spec_pad(spec, o, d, dim=-3)
    out5 = torch.zeros_like(out1)
    for i in range(t):
        out5[:, :, i] = dfop(
            spec_padded[:, :, i : i + o], coefs[:, i].unsqueeze(1), alpha[:, i].unsqueeze(1)
        )
    torch.testing.assert_allclose(out1, out5)
    # Forward method that does the padding/lookahead handling using an internal hidden state.
    dfop.freq_bins = F
    dfop.set_forward("real_hidden_state_loop")
    out6 = dfop(spec, coefs, alpha)
    torch.testing.assert_allclose(out1, out6)
//  [python]
def test_dfop():
    from df.config import config

    config.use_defaults()
    p = ModelParams()
    f = p.nb_df
    F = f * 2
    o = p.df_order
    d = p.df_lookahead
    t = 100
    spec = torch.randn(1, 1, t, F, 2)
    coefs = torch.randn(1, t, o, f, 2)
    alpha = torch.randn(1, t, 1)
    dfop = DfOp(df_bins=p.nb_df)
    dfop.set_forward("real_loop")
    out1 = dfop(spec, coefs, alpha)
    dfop.set_forward("real_strided")
    out2 = dfop(spec, coefs, alpha)
    dfop.set_forward("real_unfold")
    out3 = dfop(spec, coefs, alpha)
    dfop.set_forward("complex_strided")
    out4 = dfop(spec, coefs, alpha)
    torch.testing.assert_close(out1, out2)
    torch.testing.assert_close(out1, out3)
    torch.testing.assert_close(out1, out4)
    # This forward method requires external padding/lookahead as well as spectrogram buffer
    # handling, i.e. via a ring buffer. Could be used in real time usage.
    dfop.set_forward("real_one_step")
    spec_padded = spec_pad(spec, o, d, dim=-3)
    out5 = torch.zeros_like(out1)
    for i in range(t):
        out5[:, :, i] = dfop(
            spec_padded[:, :, i : i + o], coefs[:, i].unsqueeze(1), alpha[:, i].unsqueeze(1)
        )
    torch.testing.assert_close(out1, out5)
    # Forward method that does the padding/lookahead handling using an internal hidden state.
    dfop.freq_bins = F
    dfop.set_forward("real_hidden_state_loop")
    out6 = dfop(spec, coefs, alpha)
    torch.testing.assert_close(out1, out6)
//  [python]
def test_layer_norm_op_jit(self, X, eps, elementwise_affine, gc, dc):
        @torch.jit.script
        def jit_layer_norm(
                X: torch.Tensor,
                gamma: Optional[torch.Tensor] = None,
                beta: Optional[torch.Tensor] = None,
                axis: int = 1,
                eps: float = 1e-5,
                elementwise_affine: bool = False,
        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
            return torch.ops._caffe2.LayerNorm(
                X, gamma, beta, axis, eps, elementwise_affine)

        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.tensor(X), torch.tensor(gamma), torch.tensor(beta),
                axis, eps, elementwise_affine)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.tensor(X), None, None, axis, eps, elementwise_affine)

        torch.testing.assert_allclose(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_allclose(expected_mean, actual_mean)
        torch.testing.assert_allclose(expected_std, actual_std)
//  [python]
def test_layer_norm_op_jit(self, X, eps, elementwise_affine, gc, dc):
        @torch.jit.script
        def jit_layer_norm(
                X: torch.Tensor,
                gamma: Optional[torch.Tensor] = None,
                beta: Optional[torch.Tensor] = None,
                axis: int = 1,
                eps: float = 1e-5,
                elementwise_affine: bool = False,
        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
            return torch.ops._caffe2.LayerNorm(
                X, gamma, beta, axis, eps, elementwise_affine)

        axis = np.random.randint(0, len(X.shape))

        if elementwise_affine:
            gamma = np.random.randn(*X.shape[axis:]).astype(np.float32)
            beta = np.random.randn(*X.shape[axis:]).astype(np.float32)
            expected_norm, expected_mean, expected_std = \
                _layer_norm_with_affine_ref(axis, eps, X, gamma, beta)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.tensor(X), torch.tensor(gamma), torch.tensor(beta),
                axis, eps, elementwise_affine)
        else:
            expected_norm, expected_mean, expected_std = _layer_norm_ref(
                axis, eps, X)
            actual_norm, actual_mean, actual_std = jit_layer_norm(
                torch.tensor(X), None, None, axis, eps, elementwise_affine)

        torch.testing.assert_close(
            expected_norm, actual_norm, rtol=1e-4, atol=1e-4)
        torch.testing.assert_close(expected_mean, actual_mean)
        torch.testing.assert_close(expected_std, actual_std)
//  [python]
def test_load_mobile_module_with_debug_info(self):
        class MyTestModule(torch.nn.Module):
            def __init__(self):
                super(MyTestModule, self).__init__()

            def forward(self, x):
                return x + 5

        input = torch.tensor([3])

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter(_save_mobile_debug_info=True))
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        mobile_module_result = mobile_module(input)
        torch.testing.assert_allclose(script_module_result, mobile_module_result)

        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_allclose(script_module_result, mobile_module_forward_result)

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        torch.testing.assert_allclose(script_module_result, mobile_module_run_method_result)
//  [python]
def test_load_mobile_module_with_debug_info(self):
        class MyTestModule(torch.nn.Module):
            def __init__(self):
                super(MyTestModule, self).__init__()

            def forward(self, x):
                return x + 5

        input = torch.tensor([3])

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter(_save_mobile_debug_info=True))
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        mobile_module_result = mobile_module(input)
        torch.testing.assert_close(script_module_result, mobile_module_result)

        mobile_module_forward_result = mobile_module.forward(input)
        torch.testing.assert_close(script_module_result, mobile_module_forward_result)

        mobile_module_run_method_result = mobile_module.run_method("forward", input)
        torch.testing.assert_close(script_module_result, mobile_module_run_method_result)
//  [python]
def test_forward(self):
        x = np.random.randint(self.num_states, size=self.num_nodes)
        adj = nx.to_numpy_array(self.g)
        ref_l = np.array([adj @ (x == i) for i in range(self.num_states)])
        l = self.propagator.forward(x, self.edge_index)
        l = l.cpu().numpy()
        torch.testing.assert_allclose(l, ref_l)
//  [python]
def test_forward(self):
        x = np.random.randint(self.num_states, size=self.num_nodes)
        adj = nx.to_numpy_array(self.g)
        ref_l = np.array([adj @ (x == i) for i in range(self.num_states)])
        l = self.propagator.forward(x, self.edge_index)
        l = l.cpu().numpy()
        torch.testing.assert_close(l, ref_l)
//  [python]
def test_find_and_run_method(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, arg):
                return arg

        input = (torch.tensor([1]), )

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(*input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertFalse(has_bundled_inputs)

        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            script_module, [input], [])

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertTrue(has_bundled_inputs)

        bundled_inputs = mobile_module.run_method("get_all_bundled_inputs")
        mobile_module_result = mobile_module.forward(*bundled_inputs[0])
        torch.testing.assert_allclose(script_module_result, mobile_module_result)
//  [python]
def test_find_and_run_method(self):
        class MyTestModule(torch.nn.Module):
            def forward(self, arg):
                return arg

        input = (torch.tensor([1]), )

        script_module = torch.jit.script(MyTestModule())
        script_module_result = script_module(*input)

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertFalse(has_bundled_inputs)

        torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
            script_module, [input], [])

        buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
        buffer.seek(0)
        mobile_module = _load_for_lite_interpreter(buffer)

        has_bundled_inputs = mobile_module.find_method("get_all_bundled_inputs")
        self.assertTrue(has_bundled_inputs)

        bundled_inputs = mobile_module.run_method("get_all_bundled_inputs")
        mobile_module_result = mobile_module.forward(*bundled_inputs[0])
        torch.testing.assert_close(script_module_result, mobile_module_result)
//  [python]
def test_conv_layer_wrapper(conv_input, conv_weights, mode, clip_acts, per_channel_wts, conv_stats, expected_output):
    layer = torch.nn.Conv2d(conv_input.shape[1], expected_output.shape[1], conv_weights.shape[-1],
                            padding=1, bias=False)
    layer.weight.data = conv_weights

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts, activation_stats=conv_stats)

    with pytest.raises(RuntimeError):
        model(conv_input)

    model.eval()

    output = model(conv_input)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_conv_layer_wrapper(conv_input, conv_weights, mode, clip_acts, per_channel_wts, conv_stats, expected_output):
    layer = torch.nn.Conv2d(conv_input.shape[1], expected_output.shape[1], conv_weights.shape[-1],
                            padding=1, bias=False)
    layer.weight.data = conv_weights

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts, activation_stats=conv_stats)

    with pytest.raises(RuntimeError):
        model(conv_input)

    model.eval()

    output = model(conv_input)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_linear_layer_wrapper(linear_input, linear_weights, linear_bias,
                              mode, clip_acts, per_channel_wts, expected_output):
    layer = torch.nn.Linear(linear_input.shape[1], expected_output.shape[1], bias=True)
    layer.weight.data = linear_weights
    layer.bias.data = linear_bias

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts)

    with pytest.raises(RuntimeError):
        model(linear_input)

    model.eval()

    output = model(linear_input)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_linear_layer_wrapper(linear_input, linear_weights, linear_bias,
                              mode, clip_acts, per_channel_wts, expected_output):
    layer = torch.nn.Linear(linear_input.shape[1], expected_output.shape[1], bias=True)
    layer.weight.data = linear_weights
    layer.bias.data = linear_bias

    model = RangeLinearQuantParamLayerWrapper(layer, 8, 8, mode=mode, clip_acts=clip_acts,
                                              per_channel_wts=per_channel_wts)

    with pytest.raises(RuntimeError):
        model(linear_input)

    model.eval()

    output = model(linear_input)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_concat_layer_wrapper(inputs, concat_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantConcatWrapper(torch.nn.Module(), 8, mode, clip_acts, concat_stats)

    layer = distiller.modules.Concat(dim=1)

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    model = RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, concat_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_concat_layer_wrapper(inputs, concat_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantConcatWrapper(torch.nn.Module(), 8, mode, clip_acts, concat_stats)

    layer = distiller.modules.Concat(dim=1)

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    model = RangeLinearQuantConcatWrapper(layer, 8, mode, clip_acts, concat_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_eltwise_mult_layer_wrapper(inputs, eltwise_mult_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseMultWrapper(torch.nn.Module(), 8, mode, clip_acts, eltwise_mult_stats)

    layer = distiller.modules.EltwiseMult()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    model = RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, eltwise_mult_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_eltwise_mult_layer_wrapper(inputs, eltwise_mult_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseMultWrapper(torch.nn.Module(), 8, mode, clip_acts, eltwise_mult_stats)

    layer = distiller.modules.EltwiseMult()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    model = RangeLinearQuantEltwiseMultWrapper(layer, 8, mode, clip_acts, eltwise_mult_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_unravel_indices(
    raveled_indices: Tensor, shape: List[int], unraveled_coords_: Tensor
) -> None:
    """Unit test for converting raveled indices to an unraveled index."""
    unraveled_coords = unravel_index(raveled_indices, shape)
    torch.testing.assert_allclose(unraveled_coords, unraveled_coords_)
//  [python]
def test_unravel_indices(
    raveled_indices: Tensor, shape: List[int], unraveled_coords_: Tensor
) -> None:
    """Unit test for converting raveled indices to an unraveled index."""
    unraveled_coords = unravel_index(raveled_indices, shape)
    torch.testing.assert_close(unraveled_coords, unraveled_coords_)
//  [python]
def test_cdf(self):
        cdf = self.d.cdf(torch.zeros(self.shape)).float()
        torch.testing.assert_allclose(cdf, torch.zeros(self.shape))
        cdf = self.d.cdf(50*torch.ones(self.shape)).float()
        torch.testing.assert_allclose(cdf, torch.ones(self.shape))
//  [python]
def test_cdf(self):
        cdf = self.d.cdf(torch.zeros(self.shape)).float()
        torch.testing.assert_close(cdf, torch.zeros(self.shape))
        cdf = self.d.cdf(50*torch.ones(self.shape)).float()
        torch.testing.assert_close(cdf, torch.ones(self.shape))
//  [python]
def test_ravel_multi_index(
    unraveled_coords: Tensor, shape: List[int], raveled_indices_: Tensor
) -> None:
    """Unit tests for raveling a multi-index (i.e., unraveled coordinates).

    Args:
        unraveled_coords: Indexed tensor of coordinates.
        shape: Length of each dimension w.r.t. the grid.
        raveled_indices_: Expected tensor of coordinates.
    """
    raveled_indices = ravel_multi_index(unraveled_coords, shape)
    torch.testing.assert_allclose(raveled_indices, raveled_indices_)
//  [python]
def test_ravel_multi_index(
    unraveled_coords: Tensor, shape: List[int], raveled_indices_: Tensor
) -> None:
    """Unit tests for raveling a multi-index (i.e., unraveled coordinates).

    Args:
        unraveled_coords: Indexed tensor of coordinates.
        shape: Length of each dimension w.r.t. the grid.
        raveled_indices_: Expected tensor of coordinates.
    """
    raveled_indices = ravel_multi_index(unraveled_coords, shape)
    torch.testing.assert_close(raveled_indices, raveled_indices_)
//  [python]
def test_eltwise_add_layer_wrapper(inputs, eltwise_add_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseAddWrapper(torch.nn.Module(), 8, mode, clip_acts, test_eltwise_add_layer_wrapper)

    layer = distiller.modules.EltwiseAdd()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    model = RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, eltwise_add_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_allclose(output, expected_output)
//  [python]
def test_eltwise_add_layer_wrapper(inputs, eltwise_add_stats, mode, clip_acts, expected_output):
    with pytest.raises(ValueError):
        # Check exception on wrong layer type
        RangeLinearQuantEltwiseAddWrapper(torch.nn.Module(), 8, mode, clip_acts, test_eltwise_add_layer_wrapper)

    layer = distiller.modules.EltwiseAdd()

    with pytest.raises(NotImplementedError):
        # Check exception on no stats
        RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, activation_stats=None)

    model = RangeLinearQuantEltwiseAddWrapper(layer, 8, mode, clip_acts, eltwise_add_stats)
    model.eval()
    output = model(*inputs)

    torch.testing.assert_close(output, expected_output)
//  [python]
def test_mig(self):
        core_mig = C.MutualInformationGap()
        torch_mig = T.MutualInformationGap()

        for _ in range(3):
            z = np.random.randn(16, 16)
            a = np.random.randint(16, size=(16, 3))

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

            core_mig.update_state(z, a)
            torch_mig.update(ztm, atm)

        val = core_mig.compute()
        valtm = torch_mig.compute()

        np.testing.assert_allclose(val, valtm)
        torch.testing.assert_allclose(val, valtm)
//  [python]
def test_mig(self):
        core_mig = C.MutualInformationGap()
        torch_mig = T.MutualInformationGap()

        for _ in range(3):
            z = np.random.randn(16, 16)
            a = np.random.randint(16, size=(16, 3))

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

            core_mig.update_state(z, a)
            torch_mig.update(ztm, atm)

        val = core_mig.compute()
        valtm = torch_mig.compute()

        np.testing.assert_allclose(val, valtm)
        torch.testing.assert_close(val, valtm)
//  [python]
def test_inference_mask_generation_two_points_point_batch(self):
        model = SamModel.from_pretrained("facebook/sam-vit-base")
        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

        model.to(torch_device)
        model.eval()

        raw_image = prepare_image()

        input_points = torch.Tensor([[[400, 650]], [[220, 470]]]).cpu()  # fmt: skip

        input_points = input_points.unsqueeze(0)

        inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        iou_scores = outputs.iou_scores.cpu()
        self.assertTrue(iou_scores.shape == (1, 2, 3))
        torch.testing.assert_allclose(
            iou_scores, torch.tensor([[[0.9105, 0.9825, 0.9675], [0.7646, 0.7943, 0.7774]]]), atol=1e-4, rtol=1e-4
        )
//  [python]
def test_inference_mask_generation_two_points_point_batch(self):
        model = SamModel.from_pretrained("facebook/sam-vit-base")
        processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

        model.to(torch_device)
        model.eval()

        raw_image = prepare_image()

        input_points = torch.Tensor([[[400, 650]], [[220, 470]]]).cpu()  # fmt: skip

        input_points = input_points.unsqueeze(0)

        inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        iou_scores = outputs.iou_scores.cpu()
        self.assertTrue(iou_scores.shape == (1, 2, 3))
        torch.testing.assert_close(
            iou_scores, torch.tensor([[[0.9105, 0.9825, 0.9675], [0.7646, 0.7943, 0.7774]]]), atol=1e-4, rtol=1e-4
        )
//  [python]
def ort_validate(self, onnx_io, inputs, outputs, tolerate_small_mismatch=False):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
        # compute onnxruntime output prediction
        ort_inputs = dict((ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs))
        ort_outs = ort_session.run(None, ort_inputs)
        for i in range(0, len(outputs)):
            try:
                torch.testing.assert_allclose(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
            except AssertionError as error:
                if tolerate_small_mismatch:
                    self.assertIn("(0.00%)", str(error), str(error))
                else:
                    raise
//  [python]
def ort_validate(self, onnx_io, inputs, outputs, tolerate_small_mismatch=False):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue())
        # compute onnxruntime output prediction
        ort_inputs = dict((ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs))
        ort_outs = ort_session.run(None, ort_inputs)
        for i in range(0, len(outputs)):
            try:
                torch.testing.assert_close(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
            except AssertionError as error:
                if tolerate_small_mismatch:
                    self.assertIn("(0.00%)", str(error), str(error))
                else:
                    raise
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 64, 64)
        data_samples = [DataSample().set_gt_label([1, 2, 6])]
        # default
        model = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_score.shape, (10, ))

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

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

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

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

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

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

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

        predictions = model.predict(inputs, data_samples)
        assert predictions is data_samples
        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
//  [python]
def test_impl_rep(self):
        torch.manual_seed(1239)
        self.c = torch.tensor([.5])
        self.dim = 5
        self.scale = torch.tensor([0.8, 1.5]).unsqueeze(-1)
        self.scale.requires_grad=True
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        N = 10
        x = self.d.rsample(torch.Size([N])).squeeze(0)
        x.backward(torch.ones_like(x))

        u = self.d.cdf(x).squeeze(-1).detach().numpy()
        delta = 1e-4
        scale2 = self.scale.squeeze(-1).detach().numpy() + delta
        scale1 = self.scale.squeeze(-1).detach().numpy() - delta
        def fun(x, u, scale):
            res = np.abs(cdf_r(torch.tensor(x), torch.tensor(scale), self.c, self.dim).detach().numpy() - u)
            if x < 0: res += np.exp(-10*x)
            return res

        grad_approx = torch.zeros(N, self.scale.size(0))
        for i in range(N):
            for j in range(self.scale.size(0)):
                res2 = minimize_scalar(fun, args=(u[i, [j]], scale2[[j]]), tol=1e-7)
                res1 = minimize_scalar(fun, args=(u[i, [j]], scale1[[j]]), tol=1e-7)
                grad_approx[i, j] = float((res2.x - res1.x) / 2 / delta)
        torch.testing.assert_allclose(self.scale.grad.view(-1), grad_approx.sum(0).view(-1), rtol=0.05, atol=.5)
//  [python]
def test_impl_rep(self):
        torch.manual_seed(1239)
        self.c = torch.tensor([.5])
        self.dim = 5
        self.scale = torch.tensor([0.8, 1.5]).unsqueeze(-1)
        self.scale.requires_grad=True
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        N = 10
        x = self.d.rsample(torch.Size([N])).squeeze(0)
        x.backward(torch.ones_like(x))

        u = self.d.cdf(x).squeeze(-1).detach().numpy()
        delta = 1e-4
        scale2 = self.scale.squeeze(-1).detach().numpy() + delta
        scale1 = self.scale.squeeze(-1).detach().numpy() - delta
        def fun(x, u, scale):
            res = np.abs(cdf_r(torch.tensor(x), torch.tensor(scale), self.c, self.dim).detach().numpy() - u)
            if x < 0: res += np.exp(-10*x)
            return res

        grad_approx = torch.zeros(N, self.scale.size(0))
        for i in range(N):
            for j in range(self.scale.size(0)):
                res2 = minimize_scalar(fun, args=(u[i, [j]], scale2[[j]]), tol=1e-7)
                res1 = minimize_scalar(fun, args=(u[i, [j]], scale1[[j]]), tol=1e-7)
                grad_approx[i, j] = float((res2.x - res1.x) / 2 / delta)
        torch.testing.assert_close(self.scale.grad.view(-1), grad_approx.sum(0).view(-1), rtol=0.05, atol=.5)
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 64, 64)
        data_samples = [DataSample().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_score.shape, (10, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
        torch.testing.assert_allclose(data_samples[0].pred_score,
                                      predictions[0].pred_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 = [DataSample().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_score.shape, (10, ))

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

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_moments(self):
        self.c = torch.tensor([.7])
        self.dim = 5
        self.scale = torch.tensor([0.2, 0.5, 1., 1.5]).unsqueeze(-1)
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        means = self.d.mean
        stds = self.d.stddev
        true_means = torch.tensor([0.433607898056, 1.20011844755, 3.42020368377, 7.53105987861]).unsqueeze(-1)
        true_stds = torch.tensor([0.140120759414, 0.38460648107, 0.949379882501, 1.49865023413]).unsqueeze(-1)
        torch.testing.assert_allclose(means, true_means)
        torch.testing.assert_allclose(stds, true_stds)

        self.c = torch.tensor([.9])
        self.dim = 16
        self.scale = torch.tensor([.2, .5, .8, 1.5]).unsqueeze(-1)
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        means = self.d.mean
        stds = self.d.stddev
        true_means = torch.tensor([0.865320299218, 3.57019322017, 9.10736146478, 32.0180613104]).unsqueeze(-1)
        true_stds = torch.tensor([0.153446810207, 0.494167390237, 0.799998631679, 1.49999998722]).unsqueeze(-1)
        torch.testing.assert_allclose(means, true_means)
        torch.testing.assert_allclose(stds, true_stds)
//  [python]
def test_moments(self):
        self.c = torch.tensor([.7])
        self.dim = 5
        self.scale = torch.tensor([0.2, 0.5, 1., 1.5]).unsqueeze(-1)
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        means = self.d.mean
        stds = self.d.stddev
        true_means = torch.tensor([0.433607898056, 1.20011844755, 3.42020368377, 7.53105987861]).unsqueeze(-1)
        true_stds = torch.tensor([0.140120759414, 0.38460648107, 0.949379882501, 1.49865023413]).unsqueeze(-1)
        torch.testing.assert_close(means, true_means)
        torch.testing.assert_close(stds, true_stds)

        self.c = torch.tensor([.9])
        self.dim = 16
        self.scale = torch.tensor([.2, .5, .8, 1.5]).unsqueeze(-1)
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        means = self.d.mean
        stds = self.d.stddev
        true_means = torch.tensor([0.865320299218, 3.57019322017, 9.10736146478, 32.0180613104]).unsqueeze(-1)
        true_stds = torch.tensor([0.153446810207, 0.494167390237, 0.799998631679, 1.49999998722]).unsqueeze(-1)
        torch.testing.assert_close(means, true_means)
        torch.testing.assert_close(stds, true_stds)
//  [python]
def test_load_pretrained_model(self):
        """
        test load_state_dict_with_same_shape
        """
        model1 = MockModel()
        model2 = MockModel_()

        w1 = model1.state_dict()

        model2.load_state_dict_with_same_shape(w1)
        w2 = model2.state_dict()
        for k, p in w2.items():
            if "nn.2." not in k:
                torch.testing.assert_allclose(w1[k], p)
//  [python]
def test_load_pretrained_model(self):
        """
        test load_state_dict_with_same_shape
        """
        model1 = MockModel()
        model2 = MockModel_()

        w1 = model1.state_dict()

        model2.load_state_dict_with_same_shape(w1)
        w2 = model2.state_dict()
        for k, p in w2.items():
            if "nn.2." not in k:
                torch.testing.assert_close(w1[k], p)
//  [python]
def test_fps_sampling_registration(self):
        pos = torch.tensor([[0, 0, 0], [0.5, 0.5, 0], [0.4, 0.2, 0], [2, 2, 2], [-1, -2, -0.01]]).float()
        pair_ind = torch.tensor([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4]]).long()
        num_pos_pairs = 3
        new_ind = fps_sampling(pair_ind, pos, num_pos_pairs)
        new_pair_ind = pair_ind[new_ind]
        sol = torch.tensor([[0, 0], [3, 3], [4, 4]]).long()
        torch.testing.assert_allclose(new_pair_ind, sol)
//  [python]
def test_fps_sampling_registration(self):
        pos = torch.tensor([[0, 0, 0], [0.5, 0.5, 0], [0.4, 0.2, 0], [2, 2, 2], [-1, -2, -0.01]]).float()
        pair_ind = torch.tensor([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4]]).long()
        num_pos_pairs = 3
        new_ind = fps_sampling(pair_ind, pos, num_pos_pairs)
        new_pair_ind = pair_ind[new_ind]
        sol = torch.tensor([[0, 0], [3, 3], [4, 4]]).long()
        torch.testing.assert_close(new_pair_ind, sol)
//  [python]
def test_logZ(self):
        self.c = torch.tensor([.7])
        self.dim = 5
        self.scale = torch.tensor([0.2, 0.5, 1., 1.5]).unsqueeze(-1)
        self.scale.requires_grad = True
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        logZ = self.d.log_normalizer
        true_logZ = torch.tensor([-6.62831268665, -1.51213981062, 4.4095754933, 11.8648601183]).unsqueeze(-1)
        torch.testing.assert_allclose(logZ, true_logZ)

        grad_logZ_scale = grad(logZ, (self.scale), grad_outputs=(torch.ones_like(self.scale)))
        grad_logZ_scale = grad_logZ_scale[0]
        eps = 1e-3
        diff = HyperbolicRadius(self.dim, self.c, self.scale+eps).log_normalizer - HyperbolicRadius(self.dim, self.c, self.scale-eps).log_normalizer
        approx_grad_logZ_scale = diff / 2 / eps
        torch.testing.assert_allclose(grad_logZ_scale, approx_grad_logZ_scale)

        self.c = torch.tensor([.9])
        self.dim = 16
        self.scale = torch.tensor([.2, .5, .8, 1.5]).unsqueeze(-1)
        self.scale.requires_grad = True
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        logZ = self.d.log_normalizer
        true_logZ = torch.tensor([-10.829589564, 15.90421163, 55.8887896557, 219.5298998]).unsqueeze(-1)
        torch.testing.assert_allclose(logZ, true_logZ)

        grad_logZ_scale = grad(logZ, (self.scale), grad_outputs=(torch.ones_like(self.scale)))
        grad_logZ_scale = grad_logZ_scale[0]
        diff = HyperbolicRadius(self.dim, self.c, self.scale+eps).log_normalizer - HyperbolicRadius(self.dim, self.c, self.scale-eps).log_normalizer
        approx_grad_logZ_scale = diff / 2 / eps
        torch.testing.assert_allclose(grad_logZ_scale, approx_grad_logZ_scale)
//  [python]
def test_logZ(self):
        self.c = torch.tensor([.7])
        self.dim = 5
        self.scale = torch.tensor([0.2, 0.5, 1., 1.5]).unsqueeze(-1)
        self.scale.requires_grad = True
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        logZ = self.d.log_normalizer
        true_logZ = torch.tensor([-6.62831268665, -1.51213981062, 4.4095754933, 11.8648601183]).unsqueeze(-1)
        torch.testing.assert_close(logZ, true_logZ)

        grad_logZ_scale = grad(logZ, (self.scale), grad_outputs=(torch.ones_like(self.scale)))
        grad_logZ_scale = grad_logZ_scale[0]
        eps = 1e-3
        diff = HyperbolicRadius(self.dim, self.c, self.scale+eps).log_normalizer - HyperbolicRadius(self.dim, self.c, self.scale-eps).log_normalizer
        approx_grad_logZ_scale = diff / 2 / eps
        torch.testing.assert_close(grad_logZ_scale, approx_grad_logZ_scale)

        self.c = torch.tensor([.9])
        self.dim = 16
        self.scale = torch.tensor([.2, .5, .8, 1.5]).unsqueeze(-1)
        self.scale.requires_grad = True
        self.d = HyperbolicRadius(self.dim, self.c, self.scale)
        logZ = self.d.log_normalizer
        true_logZ = torch.tensor([-10.829589564, 15.90421163, 55.8887896557, 219.5298998]).unsqueeze(-1)
        torch.testing.assert_close(logZ, true_logZ)

        grad_logZ_scale = grad(logZ, (self.scale), grad_outputs=(torch.ones_like(self.scale)))
        grad_logZ_scale = grad_logZ_scale[0]
        diff = HyperbolicRadius(self.dim, self.c, self.scale+eps).log_normalizer - HyperbolicRadius(self.dim, self.c, self.scale-eps).log_normalizer
        approx_grad_logZ_scale = diff / 2 / eps
        torch.testing.assert_close(grad_logZ_scale, approx_grad_logZ_scale)
//  [python]
def test_iterative_policy_evaluation():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3

    policy = RandomPolicy(
        dim_state=(),
        dim_action=(),
        num_states=environment.num_states,
        num_actions=environment.num_actions,
    )
    value_function = iterative_policy_evaluation(policy, environment, GAMMA, eps=EPS)

    torch.testing.assert_allclose(
        value_function.table,
        torch.tensor([RANDOM_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
//  [python]
def test_iterative_policy_evaluation():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3

    policy = RandomPolicy(
        dim_state=(),
        dim_action=(),
        num_states=environment.num_states,
        num_actions=environment.num_actions,
    )
    value_function = iterative_policy_evaluation(policy, environment, GAMMA, eps=EPS)

    torch.testing.assert_close(
        value_function.table,
        torch.tensor([RANDOM_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
//  [python]
def test_linear_system_policy_evaluation():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3

    policy = RandomPolicy(
        dim_state=(),
        dim_action=(),
        num_states=environment.num_states,
        num_actions=environment.num_actions,
    )
    value_function = linear_system_policy_evaluation(policy, environment, GAMMA)

    torch.testing.assert_allclose(
        value_function.table,
        torch.tensor([RANDOM_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
//  [python]
def test_linear_system_policy_evaluation():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3

    policy = RandomPolicy(
        dim_state=(),
        dim_action=(),
        num_states=environment.num_states,
        num_actions=environment.num_actions,
    )
    value_function = linear_system_policy_evaluation(policy, environment, GAMMA)

    torch.testing.assert_close(
        value_function.table,
        torch.tensor([RANDOM_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
//  [python]
def test_policy_iteration():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = policy_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_allclose(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(environment, GAMMA, value_function, OPTIMAL_POLICY, pred_p)

    environment = EasyGridWorld(terminal_states=[22])
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = policy_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_allclose(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE_WITH_TERMINAL]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )

    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(
        environment, GAMMA, value_function, OPTIMAL_POLICY_WITH_TERMINAL, pred_p
    )
//  [python]
def test_policy_iteration():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = policy_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_close(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(environment, GAMMA, value_function, OPTIMAL_POLICY, pred_p)

    environment = EasyGridWorld(terminal_states=[22])
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = policy_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_close(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE_WITH_TERMINAL]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )

    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(
        environment, GAMMA, value_function, OPTIMAL_POLICY_WITH_TERMINAL, pred_p
    )
//  [python]
def assert_policy_equality(environment, gamma, value_function, true_opt_p, pred_opt_p):
    """Assert equality by checking Bellman operator equality."""
    for state in range(environment.num_states):
        environment.state = state
        next_state, reward, done, info = environment.step(pred_opt_p[state])
        reward = torch.tensor(reward, dtype=torch.get_default_dtype())
        pred_value = reward + gamma * value_function(torch.tensor(next_state))

        environment.state = state
        next_state, reward, done, info = environment.step(true_opt_p[state])
        reward = torch.tensor(reward, dtype=torch.get_default_dtype())

        true_value = reward + gamma * value_function(torch.tensor(next_state))
        torch.testing.assert_allclose(pred_value, true_value)
//  [python]
def assert_policy_equality(environment, gamma, value_function, true_opt_p, pred_opt_p):
    """Assert equality by checking Bellman operator equality."""
    for state in range(environment.num_states):
        environment.state = state
        next_state, reward, done, info = environment.step(pred_opt_p[state])
        reward = torch.tensor(reward, dtype=torch.get_default_dtype())
        pred_value = reward + gamma * value_function(torch.tensor(next_state))

        environment.state = state
        next_state, reward, done, info = environment.step(true_opt_p[state])
        reward = torch.tensor(reward, dtype=torch.get_default_dtype())

        true_value = reward + gamma * value_function(torch.tensor(next_state))
        torch.testing.assert_close(pred_value, true_value)
//  [python]
def check_module_params(module_1: torch.nn.Module, module_2: torch.nn.Module):
    module_2_params = dict(module_2.named_parameters())
    for name, param in module_1.named_parameters():
        assert name in module_2_params
        torch.testing.assert_allclose(module_2_params[name], param, rtol=0, atol=0)
//  [python]
def check_module_params(module_1: torch.nn.Module, module_2: torch.nn.Module):
    module_2_params = dict(module_2.named_parameters())
    for name, param in module_1.named_parameters():
        assert name in module_2_params
        torch.testing.assert_close(module_2_params[name], param, rtol=0, atol=0)
//  [python]
def test_pb_l2_norm():
    x = PaddedBatch(torch.randn(4, 12, 8), torch.LongTensor([3, 12, 8]))
    model = PBL2Norm()
    y = model(x)
    assert y.payload.size() == (4, 12, 8)
    torch.testing.assert_allclose(torch.ones(4, 12), y.payload.pow(2).sum(dim=2))
//  [python]
def test_pb_l2_norm():
    x = PaddedBatch(torch.randn(4, 12, 8), torch.LongTensor([3, 12, 8]))
    model = PBL2Norm()
    y = model(x)
    assert y.payload.size() == (4, 12, 8)
    torch.testing.assert_close(torch.ones(4, 12), y.payload.pow(2).sum(dim=2))
//  [python]
def test_value_iteration():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = value_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_allclose(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(environment, GAMMA, value_function, OPTIMAL_POLICY, pred_p)

    environment = EasyGridWorld(terminal_states=[22])
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = value_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_allclose(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE_WITH_TERMINAL]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )

    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(
        environment, GAMMA, value_function, OPTIMAL_POLICY_WITH_TERMINAL, pred_p
    )