jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_mask(self):
        inputs = torch.rand(16, 7 * 7, 96)
        attn = WindowMSA(embed_dims=96, window_size=(7, 7), num_heads=4)
        mask = torch.zeros((4, 49, 49))
        # Mask the first column
        mask[:, 0, :] = -100
        mask[:, :, 0] = -100
        outs = attn(inputs, mask=mask)
        inputs[:, 0, :].normal_()
        outs_with_mask = attn(inputs, mask=mask)
        torch.testing.assert_close(outs[:, 1:, :], outs_with_mask[:, 1:, :])
//  [python]
def test_qtopk(self, X, k, dim, largest, sorted):
        X, (scale, zero_point, torch_type) = X
        qX = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type)
        assume(dim < X.ndim)
        assume(k < X.shape[dim])

        unquantized_out = torch.topk(qX.dequantize(), k, dim=dim, largest=largest, sorted=sorted)

        values = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type)
        indices = torch.tensor(torch.from_numpy(X)).long()

        quantized_out = torch.topk(qX, k, dim=dim, largest=largest, sorted=sorted)

        assert(len(unquantized_out) == len(quantized_out))
        torch.testing.assert_allclose(quantized_out[0].dequantize(), unquantized_out[0])
        torch.testing.assert_allclose(quantized_out[1], unquantized_out[1])
//  [python]
def test_qtopk(self, X, k, dim, largest, sorted):
        X, (scale, zero_point, torch_type) = X
        qX = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type)
        assume(dim < X.ndim)
        assume(k < X.shape[dim])

        unquantized_out = torch.topk(qX.dequantize(), k, dim=dim, largest=largest, sorted=sorted)

        values = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type)
        indices = torch.tensor(torch.from_numpy(X)).long()

        quantized_out = torch.topk(qX, k, dim=dim, largest=largest, sorted=sorted)

        assert(len(unquantized_out) == len(quantized_out))
        torch.testing.assert_close(quantized_out[0].dequantize(), unquantized_out[0])
        torch.testing.assert_close(quantized_out[1], unquantized_out[1])
//  [python]
def test_qtopk_nhwc(self, X, k, dim, largest, sorted):
        # X is NHWC, we permute to view as NCHW but keep NHWC in memory
        X, (scale, zero_point, torch_type) = X
        qX = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type).permute([0, 3, 1, 2])
        X = np.transpose(X, [0, 3, 1, 2])
        assume(dim < X.ndim)
        assume(k < X.shape[dim])

        unquantized_out = torch.topk(qX.dequantize(), k, dim=dim, largest=largest, sorted=sorted)

        values = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type)
        indices = torch.tensor(torch.from_numpy(X)).long()

        quantized_out = torch.topk(qX, k, dim=dim, largest=largest, sorted=sorted)

        assert(len(unquantized_out) == len(quantized_out))
        torch.testing.assert_allclose(quantized_out[0].dequantize(), unquantized_out[0])
        torch.testing.assert_allclose(quantized_out[1], unquantized_out[1])
//  [python]
def test_qtopk_nhwc(self, X, k, dim, largest, sorted):
        # X is NHWC, we permute to view as NCHW but keep NHWC in memory
        X, (scale, zero_point, torch_type) = X
        qX = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type).permute([0, 3, 1, 2])
        X = np.transpose(X, [0, 3, 1, 2])
        assume(dim < X.ndim)
        assume(k < X.shape[dim])

        unquantized_out = torch.topk(qX.dequantize(), k, dim=dim, largest=largest, sorted=sorted)

        values = torch.quantize_per_tensor(torch.from_numpy(X), scale, zero_point, torch_type)
        indices = torch.tensor(torch.from_numpy(X)).long()

        quantized_out = torch.topk(qX, k, dim=dim, largest=largest, sorted=sorted)

        assert(len(unquantized_out) == len(quantized_out))
        torch.testing.assert_close(quantized_out[0].dequantize(), unquantized_out[0])
        torch.testing.assert_close(quantized_out[1], unquantized_out[1])
//  [python]
def compare_eps_attrib_values(c: Cumulant):
    for non_centered in [False, True]:
        to_expand = py_to_rust(eps_attrib(c, non_centered=non_centered))
        r_to_expand = r_eps_attrib(py_to_rust(c).cast_cumulant(), non_centered=non_centered)

        sample_spec = RunDiscreteVarAllSpec.create_full_from_circuits(to_expand)
        sampled = Sampler(sample_spec)(to_expand).evaluate()
        r_sample_spec = RunDiscreteVarAllSpec.create_full_from_circuits(r_to_expand)
        # TODO: make sampling work?
        r_sampled = Sampler(r_sample_spec)(r_to_expand.cast_module().substitute()).evaluate()

        atol = 1e-6
        rtol = 1e-5
        torch.testing.assert_allclose(sampled, r_sampled, atol=atol, rtol=rtol)
//  [python]
def compare_eps_attrib_values(c: Cumulant):
    for non_centered in [False, True]:
        to_expand = py_to_rust(eps_attrib(c, non_centered=non_centered))
        r_to_expand = r_eps_attrib(py_to_rust(c).cast_cumulant(), non_centered=non_centered)

        sample_spec = RunDiscreteVarAllSpec.create_full_from_circuits(to_expand)
        sampled = Sampler(sample_spec)(to_expand).evaluate()
        r_sample_spec = RunDiscreteVarAllSpec.create_full_from_circuits(r_to_expand)
        # TODO: make sampling work?
        r_sampled = Sampler(r_sample_spec)(r_to_expand.cast_module().substitute()).evaluate()

        atol = 1e-6
        rtol = 1e-5
        torch.testing.assert_close(sampled, r_sampled, atol=atol, rtol=rtol)
//  [python]
def test_cat_nhwc(self, X, relu):
        # X is NHWC
        X, (scale, zero_point, torch_type) = X

        # Tile out X so # channels is > 64
        X = np.repeat(X, 70 / X.shape[3], 3)
        X = torch.from_numpy(np.ascontiguousarray(X))
        Y = X.clone()
        Y = torch.from_numpy(np.ascontiguousarray(Y))
        # Here, we quantize and get quantized tensors in NHWC for both dims and strides. The
        # permute switches it so that the tensor looks like NCHW but it laid out in memory as
        # NHWC.
        qX = torch.quantize_per_tensor(X, scale, zero_point, torch_type).permute([0, 3, 1, 2])
        qY = torch.quantize_per_tensor(Y, scale, zero_point, torch_type).permute([0, 3, 1, 2])

        ref = torch.cat([qX.dequantize(), qY.dequantize()], dim=1)
        if relu:
            ref[ref < 0] = 0.0
        ref = torch.quantize_per_tensor(ref, scale=scale, zero_point=zero_point, dtype=torch_type)

        if relu:
            out = torch.ops.quantized.cat_relu(
                [qX, qY], dim=1, scale=scale, zero_point=zero_point)
        else:
            out = torch.ops.quantized.cat([qX, qY], dim=1, scale=scale, zero_point=zero_point)

        torch.testing.assert_allclose(out.dequantize(), ref.dequantize())
        self.assertNotEqual(out.stride(), sorted(out.stride()))
//  [python]
def test_cat_nhwc(self, X, relu):
        # X is NHWC
        X, (scale, zero_point, torch_type) = X

        # Tile out X so # channels is > 64
        X = np.repeat(X, 70 / X.shape[3], 3)
        X = torch.from_numpy(np.ascontiguousarray(X))
        Y = X.clone()
        Y = torch.from_numpy(np.ascontiguousarray(Y))
        # Here, we quantize and get quantized tensors in NHWC for both dims and strides. The
        # permute switches it so that the tensor looks like NCHW but it laid out in memory as
        # NHWC.
        qX = torch.quantize_per_tensor(X, scale, zero_point, torch_type).permute([0, 3, 1, 2])
        qY = torch.quantize_per_tensor(Y, scale, zero_point, torch_type).permute([0, 3, 1, 2])

        ref = torch.cat([qX.dequantize(), qY.dequantize()], dim=1)
        if relu:
            ref[ref < 0] = 0.0
        ref = torch.quantize_per_tensor(ref, scale=scale, zero_point=zero_point, dtype=torch_type)

        if relu:
            out = torch.ops.quantized.cat_relu(
                [qX, qY], dim=1, scale=scale, zero_point=zero_point)
        else:
            out = torch.ops.quantized.cat([qX, qY], dim=1, scale=scale, zero_point=zero_point)

        torch.testing.assert_close(out.dequantize(), ref.dequantize())
        self.assertNotEqual(out.stride(), sorted(out.stride()))
//  [python]
def test_embedding_byte(self, num_embeddings, embedding_dim):
        quant_op = torch.ops.quantized.embedding_byte
        prepack_op = torch.ops.quantized.embedding_bag_prepack

        weights = torch.from_numpy((np.random.random_sample((
            num_embeddings, embedding_dim)) + 1).astype(np.float32))

        obs = PerChannelMinMaxObserver(dtype=torch.quint8, qscheme=torch.per_channel_affine_float_qparams, ch_axis=0)
        obs(weights)
        # Get the scale and zero point for the weight tensor
        qparams = obs.calculate_qparams()

        # Quantize the weights to 8bits
        qweight = torch.quantize_per_channel(weights, qparams[0], qparams[1], axis=0, dtype=torch.quint8)
        max_segments = 5
        max_segment_length = 20
        num_lengths = np.random.randint(1, max_segments + 1)
        lengths = np.random.randint(1, max_segment_length + 1,
                                    size=num_lengths).astype(np.int32)
        num_indices = np.sum(lengths)
        indices = torch.from_numpy(np.random.randint(
            low=0, high=num_embeddings, size=num_indices, dtype=np.int64))

        packed_weight = prepack_op(qweight)
        qresult = quant_op(packed_weight, indices, sparse=False)

        ref = torch.embedding(weights, indices, padding_idx=-1, scale_grad_by_freq=False, sparse=False)
        torch.testing.assert_allclose(ref, qresult, atol=0.005, rtol=1e-3)
//  [python]
def test_embedding_byte(self, num_embeddings, embedding_dim):
        quant_op = torch.ops.quantized.embedding_byte
        prepack_op = torch.ops.quantized.embedding_bag_prepack

        weights = torch.from_numpy((np.random.random_sample((
            num_embeddings, embedding_dim)) + 1).astype(np.float32))

        obs = PerChannelMinMaxObserver(dtype=torch.quint8, qscheme=torch.per_channel_affine_float_qparams, ch_axis=0)
        obs(weights)
        # Get the scale and zero point for the weight tensor
        qparams = obs.calculate_qparams()

        # Quantize the weights to 8bits
        qweight = torch.quantize_per_channel(weights, qparams[0], qparams[1], axis=0, dtype=torch.quint8)
        max_segments = 5
        max_segment_length = 20
        num_lengths = np.random.randint(1, max_segments + 1)
        lengths = np.random.randint(1, max_segment_length + 1,
                                    size=num_lengths).astype(np.int32)
        num_indices = np.sum(lengths)
        indices = torch.from_numpy(np.random.randint(
            low=0, high=num_embeddings, size=num_indices, dtype=np.int64))

        packed_weight = prepack_op(qweight)
        qresult = quant_op(packed_weight, indices, sparse=False)

        ref = torch.embedding(weights, indices, padding_idx=-1, scale_grad_by_freq=False, sparse=False)
        torch.testing.assert_close(ref, qresult, atol=0.005, rtol=1e-3)
//  [python]
def test_crop(self):
        input = torch.rand(64, 256, 13, 24)
        other = torch.rand(32, 8, 6)
        output = ops.crop_or_pad_as(input, other)
        self.assertEqual(output.shape[:-2], input.shape[:-2])
        self.assertEqual(output.shape[-2:], other.shape[-2:])
        torch.testing.assert_allclose(
            output, input[..., 2:10, 9:15])
//  [python]
def test_crop(self):
        input = torch.rand(64, 256, 13, 24)
        other = torch.rand(32, 8, 6)
        output = ops.crop_or_pad_as(input, other)
        self.assertEqual(output.shape[:-2], input.shape[:-2])
        self.assertEqual(output.shape[-2:], other.shape[-2:])
        torch.testing.assert_close(
            output, input[..., 2:10, 9:15])
//  [python]
def test_pad(self):
        input = torch.rand(64, 256, 7, 9)
        other = torch.rand(32, 18, 19)
        output = ops.crop_or_pad_as(input, other, pad_val=0.5)
        self.assertEqual(output.shape[:-2], input.shape[:-2])
        self.assertEqual(output.shape[-2:], other.shape[-2:])
        torch.testing.assert_allclose(
            output[..., 5:12, 5:14], input)
        self.assertEqual(output[0, 0, 0, 0], 0.5)
//  [python]
def test_pad(self):
        input = torch.rand(64, 256, 7, 9)
        other = torch.rand(32, 18, 19)
        output = ops.crop_or_pad_as(input, other, pad_val=0.5)
        self.assertEqual(output.shape[:-2], input.shape[:-2])
        self.assertEqual(output.shape[-2:], other.shape[-2:])
        torch.testing.assert_close(
            output[..., 5:12, 5:14], input)
        self.assertEqual(output[0, 0, 0, 0], 0.5)
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

        class M(torch.nn.Module):
            def __init__(self):
                super(M, self).__init__()
                self.weight_0 = torch.nn.Parameter(torch.Tensor(torch.rand(weight_0_shape)))
                self.weight_1 = torch.nn.Parameter(torch.Tensor(torch.rand(weight_1_shape)))

            def forward(self, x):
                o = F.linear(x, self.weight_0)
                o = F.dropout(o, training=self.training)
                o = F.linear(o, self.weight_1)
                return o

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        torch.testing.assert_allclose(ref_res, res, rtol=1e-2, atol=1e-3)
//  [python]
def test_remove_dropout(self):
        weight_0_shape = (20, 5)
        weight_1_shape = (20, 20)
        input_shape = (10, 5)

        class M(torch.nn.Module):
            def __init__(self):
                super(M, self).__init__()
                self.weight_0 = torch.nn.Parameter(torch.Tensor(torch.rand(weight_0_shape)))
                self.weight_1 = torch.nn.Parameter(torch.Tensor(torch.rand(weight_1_shape)))

            def forward(self, x):
                o = F.linear(x, self.weight_0)
                o = F.dropout(o, training=self.training)
                o = F.linear(o, self.weight_1)
                return o

        data = torch.rand(input_shape)
        m = M()
        m = torch.jit.script(m)
        with self.assertRaisesRegex(RuntimeError, r'Dropout removal module in training mode is not yet supported'):
            torch._C._jit_pass_remove_dropout(m._c)
        m.eval()
        ref_res = m(data)
        # Need to inline otherwise we see instances of Function.
        # We would have to use torch.linear/dropout to get around it otherwise.
        from torch.jit._recursive import wrap_cpp_module
        m = wrap_cpp_module(torch._C._freeze_module(m._c))
        torch._C._jit_pass_remove_dropout(m._c)
        res = m(data)
        FileCheck().check_not("aten::dropout").run(str(m.graph))
        torch.testing.assert_close(ref_res, res, rtol=1e-2, atol=1e-3)
//  [python]
def test_crop_and_pad(self):
        input = torch.rand(64, 256, 20, 9)
        other = torch.rand(32, 18, 19)
        output = ops.crop_or_pad_as(input, other, pad_val=0.5)
        self.assertEqual(output.shape[:-2], input.shape[:-2])
        self.assertEqual(output.shape[-2:], other.shape[-2:])
        torch.testing.assert_allclose(
            output[..., :, 5:14], input[..., 1:19, :])
        self.assertEqual(output[0, 0, 0, 0], 0.5)
//  [python]
def test_crop_and_pad(self):
        input = torch.rand(64, 256, 20, 9)
        other = torch.rand(32, 18, 19)
        output = ops.crop_or_pad_as(input, other, pad_val=0.5)
        self.assertEqual(output.shape[:-2], input.shape[:-2])
        self.assertEqual(output.shape[-2:], other.shape[-2:])
        torch.testing.assert_close(
            output[..., :, 5:14], input[..., 1:19, :])
        self.assertEqual(output[0, 0, 0, 0], 0.5)
//  [python]
def test_expected_returns_standardized(self):
        """
        Test the expected return method for standardized data.

        Notes
        -----
        Test that the expected returns are correct.
        """
        discounted_returns = self.loss.compute_true_value_function(rewards=self.rewards)
        torch.testing.assert_allclose(torch.mean(discounted_returns), 0.0)
        torch.testing.assert_allclose(torch.std(discounted_returns), 1.0)
//  [python]
def test_expected_returns_standardized(self):
        """
        Test the expected return method for standardized data.

        Notes
        -----
        Test that the expected returns are correct.
        """
        discounted_returns = self.loss.compute_true_value_function(rewards=self.rewards)
        torch.testing.assert_close(torch.mean(discounted_returns), 0.0)
        torch.testing.assert_close(torch.std(discounted_returns), 1.0)
//  [python]
def test_true_values(self):
        """
        Compute the discounted returns of the actual rewards.

        Notes
        -------
        Checks the actual values return in the discounted returns without
        standardization and with the simple case of gamma=1
        """
        discounted_return = torch.tensor([15, 14, 12, 9, 5])
        value_function = self.loss.compute_true_value_function(
            rewards=self.rewards, gamma=1, standardize=False
        )
        torch.testing.assert_allclose(discounted_return, value_function)
//  [python]
def test_true_values(self):
        """
        Compute the discounted returns of the actual rewards.

        Notes
        -------
        Checks the actual values return in the discounted returns without
        standardization and with the simple case of gamma=1
        """
        discounted_return = torch.tensor([15, 14, 12, 9, 5])
        value_function = self.loss.compute_true_value_function(
            rewards=self.rewards, gamma=1, standardize=False
        )
        torch.testing.assert_close(discounted_return, value_function)
//  [python]
def assert_arrays_almost_equals(a, b):
    if a is None or b is None:
        assert a == b
    else:
        if torch.is_tensor(a) or torch.is_tensor(b):
            torch.testing.assert_allclose(a, b)
        elif isinstance(a, list) or isinstance(b, list):
            assert a == b
        else:
            assert np.allclose(a, b)
//  [python]
def assert_arrays_almost_equals(a, b):
    if a is None or b is None:
        assert a == b
    else:
        if torch.is_tensor(a) or torch.is_tensor(b):
            torch.testing.assert_close(a, b)
        elif isinstance(a, list) or isinstance(b, list):
            assert a == b
        else:
            assert np.allclose(a, b)
//  [python]
def test_normal_cdf(self):
        from scipy.stats import norm

        scheduler = DDPMScheduler(num_train_timesteps=10)
        inferer = DiffusionInferer(scheduler=scheduler)

        x = torch.linspace(-10, 10, 20)
        cdf_approx = inferer._approx_standard_normal_cdf(x)
        cdf_true = norm.cdf(x)
        torch.testing.assert_allclose(cdf_approx, cdf_true, atol=1e-3, rtol=1e-5)
//  [python]
def test_normal_cdf(self):
        from scipy.stats import norm

        scheduler = DDPMScheduler(num_train_timesteps=10)
        inferer = DiffusionInferer(scheduler=scheduler)

        x = torch.linspace(-10, 10, 20)
        cdf_approx = inferer._approx_standard_normal_cdf(x)
        cdf_true = norm.cdf(x)
        torch.testing.assert_close(cdf_approx, cdf_true, atol=1e-3, rtol=1e-5)
//  [python]
def test_1d(self):
        """Test jacobian function for 1D inputs."""
        x = torch.ones(1, requires_grad=True)
        f = 2 * x

        jac = compute_jacobian(f, x)
        torch.testing.assert_allclose(jac[0, 0], 2)
//  [python]
def test_1d(self):
        """Test jacobian function for 1D inputs."""
        x = torch.ones(1, requires_grad=True)
        f = 2 * x

        jac = compute_jacobian(f, x)
        torch.testing.assert_close(jac[0, 0], 2)
//  [python]
def test_0d(self):
        x = torch.ones(2, 2, requires_grad=True)
        A = torch.tensor([[1., 2.], [3., 4.]])
        f = A * x
        f = torch.sum(f)

        jac = compute_jacobian(f, x)
        torch.testing.assert_allclose(jac, A)
//  [python]
def test_0d(self):
        x = torch.ones(2, 2, requires_grad=True)
        A = torch.tensor([[1., 2.], [3., 4.]])
        f = A * x
        f = torch.sum(f)

        jac = compute_jacobian(f, x)
        torch.testing.assert_close(jac, A)
//  [python]
a = torch.load(sys.argv[1])
b = torch.load(sys.argv[2])
torch.testing.assert_allclose(a,b, rtol=0.01, atol=0.01)
//  [python]
a = torch.load(sys.argv[1])
b = torch.load(sys.argv[2])
torch.testing.assert_close(a,b, rtol=0.01, atol=0.01)
//  [python]
def test_2d(self):
        """Test jacobian computation."""
        x = torch.ones(2, 1, requires_grad=True)
        A = torch.tensor([[1., 2.], [3., 4.]])
        f = A @ x

        jac = compute_jacobian(f, x)
        torch.testing.assert_allclose(A, jac[:, 0, :, 0])

        # Test both multiple runs
        jac = compute_jacobian(f.squeeze(-1), x)
        torch.testing.assert_allclose(A, jac.squeeze(-1))
//  [python]
def test_2d(self):
        """Test jacobian computation."""
        x = torch.ones(2, 1, requires_grad=True)
        A = torch.tensor([[1., 2.], [3., 4.]])
        f = A @ x

        jac = compute_jacobian(f, x)
        torch.testing.assert_close(A, jac[:, 0, :, 0])

        # Test both multiple runs
        jac = compute_jacobian(f.squeeze(-1), x)
        torch.testing.assert_close(A, jac.squeeze(-1))
//  [python]
def test_2d_output(self):
        """Test jacobian with 2d input and output"""
        x = torch.ones(2, 2, requires_grad=True)
        A = torch.tensor([[1., 2.], [3., 4.]])
        f = A * x

        jac = compute_jacobian(f, x)
        torch.testing.assert_allclose(jac.shape, 2)
        torch.testing.assert_allclose(jac.sum(dim=0).sum(dim=0), A)
//  [python]
def test_2d_output(self):
        """Test jacobian with 2d input and output"""
        x = torch.ones(2, 2, requires_grad=True)
        A = torch.tensor([[1., 2.], [3., 4.]])
        f = A * x

        jac = compute_jacobian(f, x)
        torch.testing.assert_close(jac.shape, 2)
        torch.testing.assert_close(jac.sum(dim=0).sum(dim=0), A)
//  [python]
def test_empty_module(self):
        """Test module without any param."""
        module = self._get_empty_module()
        in_data = torch.rand(1)
        ref_out = module(in_data)
        module = FlattenParamsWrapper(module)
        assert len(list(module.parameters())) == 0
        assert len(module.state_dict()) == 0
        fpw_out = module(in_data)
        torch.testing.assert_allclose(ref_out, fpw_out)
//  [python]
def test_empty_module(self):
        """Test module without any param."""
        module = self._get_empty_module()
        in_data = torch.rand(1)
        ref_out = module(in_data)
        module = FlattenParamsWrapper(module)
        assert len(list(module.parameters())) == 0
        assert len(module.state_dict()) == 0
        fpw_out = module(in_data)
        torch.testing.assert_close(ref_out, fpw_out)
//  [python]
def test_independent_implicit_quantile() -> None:
    num_samples = NUM_SAMPLES

    # # Normal distrib
    distr_mean = torch.Tensor([10.0])
    distr_std = torch.Tensor([4.0])
    distr_pp10 = distr_mean - 1.282 * distr_std
    distr_pp90 = distr_mean + 1.282 * distr_std
    distr = Normal(loc=distr_mean, scale=distr_std)

    samples = distr.sample((num_samples,))
    learned_mean, learned_std, learned_pp10, learned_pp90 = learn_distribution(
        ImplicitQuantileOutput(output_domain="Real"),
        samples=samples,
        num_epochs=50,
        learning_rate=1e-2,
    )

    torch.testing.assert_allclose(
        learned_mean, distr_mean.squeeze(), rtol=0.1, atol=0.1 * 10
    )
    torch.testing.assert_allclose(
        learned_std, distr_std.squeeze(), rtol=0.1, atol=0.1 * 4
    )
    torch.testing.assert_allclose(
        learned_pp90, distr_pp90.squeeze(), rtol=0.1, atol=0.1 * 4
    )
    torch.testing.assert_allclose(
        learned_pp10, distr_pp10.squeeze(), rtol=0.1, atol=0.1 * 4
    )

    # Uniform distrib
    a = torch.Tensor([0.0])
    b = torch.Tensor([20.0])
    distr_mean = 0.5 * (a + b)
    distr_std = (1.0 / 12.0 * (b - a) ** 2) ** 0.5
    distr_pp10 = 0.1 * (a + b)
    distr_pp90 = 0.9 * (a + b)
    distr = Uniform(low=a, high=b)

    samples = distr.sample((num_samples,))
    learned_mean, learned_std, learned_pp10, learned_pp90 = learn_distribution(
        ImplicitQuantileOutput(output_domain="Positive"),
        samples=samples,
        num_epochs=50,
        learning_rate=1e-2,
    )

    torch.testing.assert_allclose(
        learned_mean, distr_mean.squeeze(), atol=1.0, rtol=0.1
    )
    torch.testing.assert_allclose(learned_std, distr_std.squeeze(), atol=0.5, rtol=0.1)
    torch.testing.assert_allclose(
        learned_pp90, distr_pp90.squeeze(), rtol=0.1, atol=0.1 * 18
    )
    torch.testing.assert_allclose(
        learned_pp10, distr_pp10.squeeze(), rtol=0.2, atol=0.2 * 2
    )

    # Bernoulli distrib
    distr_mean = torch.Tensor([0.2])
    distr_std = distr_mean * (1 - distr_mean)
    distr_pp10 = torch.Tensor([0.0])
    distr_pp90 = torch.Tensor([1.0])
    distr = Bernoulli(probs=distr_mean)

    samples = distr.sample((num_samples,))
    learned_mean, learned_std, learned_pp10, learned_pp90 = learn_distribution(
        ImplicitQuantileOutput(output_domain="Positive"),
        samples=samples,
        num_epochs=50,
        learning_rate=1e-2,
    )

    torch.testing.assert_allclose(
        learned_mean, distr_mean.squeeze(), atol=1.0, rtol=0.1
    )
    torch.testing.assert_allclose(learned_std, distr_std.squeeze(), atol=0.5, rtol=0.1)
    torch.testing.assert_allclose(
        learned_pp90, distr_pp90.squeeze(), rtol=0.1, atol=0.1 * 18
    )
    torch.testing.assert_allclose(
        learned_pp10, distr_pp10.squeeze(), rtol=0.1, atol=0.1 * 2
    )
//  [python]
def test_independent_implicit_quantile() -> None:
    num_samples = NUM_SAMPLES

    # # Normal distrib
    distr_mean = torch.Tensor([10.0])
    distr_std = torch.Tensor([4.0])
    distr_pp10 = distr_mean - 1.282 * distr_std
    distr_pp90 = distr_mean + 1.282 * distr_std
    distr = Normal(loc=distr_mean, scale=distr_std)

    samples = distr.sample((num_samples,))
    learned_mean, learned_std, learned_pp10, learned_pp90 = learn_distribution(
        ImplicitQuantileOutput(output_domain="Real"),
        samples=samples,
        num_epochs=50,
        learning_rate=1e-2,
    )

    torch.testing.assert_close(
        learned_mean, distr_mean.squeeze(), rtol=0.1, atol=0.1 * 10
    )
    torch.testing.assert_close(
        learned_std, distr_std.squeeze(), rtol=0.1, atol=0.1 * 4
    )
    torch.testing.assert_close(
        learned_pp90, distr_pp90.squeeze(), rtol=0.1, atol=0.1 * 4
    )
    torch.testing.assert_close(
        learned_pp10, distr_pp10.squeeze(), rtol=0.1, atol=0.1 * 4
    )

    # Uniform distrib
    a = torch.Tensor([0.0])
    b = torch.Tensor([20.0])
    distr_mean = 0.5 * (a + b)
    distr_std = (1.0 / 12.0 * (b - a) ** 2) ** 0.5
    distr_pp10 = 0.1 * (a + b)
    distr_pp90 = 0.9 * (a + b)
    distr = Uniform(low=a, high=b)

    samples = distr.sample((num_samples,))
    learned_mean, learned_std, learned_pp10, learned_pp90 = learn_distribution(
        ImplicitQuantileOutput(output_domain="Positive"),
        samples=samples,
        num_epochs=50,
        learning_rate=1e-2,
    )

    torch.testing.assert_close(
        learned_mean, distr_mean.squeeze(), atol=1.0, rtol=0.1
    )
    torch.testing.assert_close(learned_std, distr_std.squeeze(), atol=0.5, rtol=0.1)
    torch.testing.assert_close(
        learned_pp90, distr_pp90.squeeze(), rtol=0.1, atol=0.1 * 18
    )
    torch.testing.assert_close(
        learned_pp10, distr_pp10.squeeze(), rtol=0.2, atol=0.2 * 2
    )

    # Bernoulli distrib
    distr_mean = torch.Tensor([0.2])
    distr_std = distr_mean * (1 - distr_mean)
    distr_pp10 = torch.Tensor([0.0])
    distr_pp90 = torch.Tensor([1.0])
    distr = Bernoulli(probs=distr_mean)

    samples = distr.sample((num_samples,))
    learned_mean, learned_std, learned_pp10, learned_pp90 = learn_distribution(
        ImplicitQuantileOutput(output_domain="Positive"),
        samples=samples,
        num_epochs=50,
        learning_rate=1e-2,
    )

    torch.testing.assert_close(
        learned_mean, distr_mean.squeeze(), atol=1.0, rtol=0.1
    )
    torch.testing.assert_close(learned_std, distr_std.squeeze(), atol=0.5, rtol=0.1)
    torch.testing.assert_close(
        learned_pp90, distr_pp90.squeeze(), rtol=0.1, atol=0.1 * 18
    )
    torch.testing.assert_close(
        learned_pp10, distr_pp10.squeeze(), rtol=0.1, atol=0.1 * 2
    )
//  [python]
def test_shared_params_pnorm_after_step(self):
        # incorrect parameter sharing is likely to cause problems after an
        # optimization step
        module = self._get_shared_params_transformer()
        ref_pnorm_after_step = self._get_pnorm_after_step(module)

        module = self._get_shared_params_transformer()  # recreate
        flat_module = FlattenParamsWrapper(module)
        flat_pnorm_after_step = self._get_pnorm_after_step(flat_module)

        torch.testing.assert_allclose(ref_pnorm_after_step, flat_pnorm_after_step)
//  [python]
def test_shared_params_pnorm_after_step(self):
        # incorrect parameter sharing is likely to cause problems after an
        # optimization step
        module = self._get_shared_params_transformer()
        ref_pnorm_after_step = self._get_pnorm_after_step(module)

        module = self._get_shared_params_transformer()  # recreate
        flat_module = FlattenParamsWrapper(module)
        flat_pnorm_after_step = self._get_pnorm_after_step(flat_module)

        torch.testing.assert_close(ref_pnorm_after_step, flat_pnorm_after_step)
//  [python]
def test_flatten_nothing(self):
        """Testing nothing is flatten case."""
        module = self._get_transformer()
        ref_out = self._compute_output(module)
        ref_state_dict = module.state_dict()
        for k, v in ref_state_dict.items():
            ref_state_dict[k] = v.clone()
        module = FlattenParamsWrapper(module, param_list=[[]])
        fpw_state_dict = module.state_dict()
        assert ref_state_dict.keys() == fpw_state_dict.keys()
        for k, v in ref_state_dict.items():
            torch.testing.assert_allclose(v, fpw_state_dict[k])
        fpw_out = self._compute_output(module)
        torch.testing.assert_allclose(ref_out, fpw_out)
//  [python]
def test_flatten_nothing(self):
        """Testing nothing is flatten case."""
        module = self._get_transformer()
        ref_out = self._compute_output(module)
        ref_state_dict = module.state_dict()
        for k, v in ref_state_dict.items():
            ref_state_dict[k] = v.clone()
        module = FlattenParamsWrapper(module, param_list=[[]])
        fpw_state_dict = module.state_dict()
        assert ref_state_dict.keys() == fpw_state_dict.keys()
        for k, v in ref_state_dict.items():
            torch.testing.assert_close(v, fpw_state_dict[k])
        fpw_out = self._compute_output(module)
        torch.testing.assert_close(ref_out, fpw_out)
//  [python]
def test_reward_to_go():
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

    terminal = torch.tensor([False, False, False])
    actual = reward_to_go(reward, value, terminal, gamma)
    torch.testing.assert_allclose(actual, torch.tensor([9., 8., 6.]))

    terminal = torch.tensor([False, True, False])
    actual = reward_to_go(reward, value, terminal, gamma)
    torch.testing.assert_allclose(actual, torch.tensor([3., 2., 6.]))
//  [python]
def test_reward_to_go():
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

    terminal = torch.tensor([False, False, False])
    actual = reward_to_go(reward, value, terminal, gamma)
    torch.testing.assert_close(actual, torch.tensor([9., 8., 6.]))

    terminal = torch.tensor([False, True, False])
    actual = reward_to_go(reward, value, terminal, gamma)
    torch.testing.assert_close(actual, torch.tensor([3., 2., 6.]))
//  [python]
def test_prescriber_interpolate():
    prescriber = Prescriber(
        prescribed_name="a", mask_name="mask", mask_value=1, interpolate=True
    )
    data = {
        "a": torch.zeros(2, 4, 4),
        "b": torch.ones(2, 4, 4) * 4.0,
        "mask": torch.ones(2, 4, 4) * 0.25,
    }
    target = {
        "a": torch.ones(2, 4, 4) * 4.0,
        "b": torch.zeros(2, 4, 4),
    }
    prescribed_gen = prescriber(data=data, gen_norm=data, target_norm=target)
    torch.testing.assert_allclose(prescribed_gen["a"], torch.ones(2, 4, 4))
    # check that the other variable is not changed
    torch.testing.assert_allclose(prescribed_gen["b"], torch.ones(2, 4, 4) * 4.0)
//  [python]
def test_prescriber_interpolate():
    prescriber = Prescriber(
        prescribed_name="a", mask_name="mask", mask_value=1, interpolate=True
    )
    data = {
        "a": torch.zeros(2, 4, 4),
        "b": torch.ones(2, 4, 4) * 4.0,
        "mask": torch.ones(2, 4, 4) * 0.25,
    }
    target = {
        "a": torch.ones(2, 4, 4) * 4.0,
        "b": torch.zeros(2, 4, 4),
    }
    prescribed_gen = prescriber(data=data, gen_norm=data, target_norm=target)
    torch.testing.assert_close(prescribed_gen["a"], torch.ones(2, 4, 4))
    # check that the other variable is not changed
    torch.testing.assert_close(prescribed_gen["b"], torch.ones(2, 4, 4) * 4.0)
//  [python]
def test_init(self, dim_reward):
        value_function = TabularValueFunction(num_states=4, dim_reward=(dim_reward,))
        torch.testing.assert_allclose(
            value_function.table, torch.zeros(1, 4, dim_reward)
        )
//  [python]
def test_init(self, dim_reward):
        value_function = TabularValueFunction(num_states=4, dim_reward=(dim_reward,))
        torch.testing.assert_close(
            value_function.table, torch.zeros(1, 4, dim_reward)
        )
//  [python]
def test_set_value(self, dim_reward):
        value_function = TabularValueFunction(num_states=4, dim_reward=(dim_reward,))
        value_function.set_value(2, torch.ones(dim_reward))
        table = torch.zeros(1, 4, dim_reward)
        table[:, 2, :] = 1.0
        torch.testing.assert_allclose(value_function.table, table)
//  [python]
def test_set_value(self, dim_reward):
        value_function = TabularValueFunction(num_states=4, dim_reward=(dim_reward,))
        value_function.set_value(2, torch.ones(dim_reward))
        table = torch.zeros(1, 4, dim_reward)
        table[:, 2, :] = 1.0
        torch.testing.assert_close(value_function.table, table)
//  [python]
def test_init(self, dim_reward):
        value_function = TabularQFunction(
            num_states=4, num_actions=2, dim_reward=(dim_reward,)
        )
        torch.testing.assert_allclose(
            value_function.table, torch.zeros(2, 4, dim_reward)
        )
//  [python]
def test_init(self, dim_reward):
        value_function = TabularQFunction(
            num_states=4, num_actions=2, dim_reward=(dim_reward,)
        )
        torch.testing.assert_close(
            value_function.table, torch.zeros(2, 4, dim_reward)
        )
//  [python]
def test_prescriber():
    """Test that the prescriber overwrites the generated data in the masked region."""
    prescriber = Prescriber(prescribed_name="a", mask_name="mask", mask_value=0)
    data = {
        "a": torch.rand(2, 4, 4),
        "b": torch.rand(2, 4, 4),
        "mask": torch.ones(2, 4, 4),
    }
    target = {
        "a": torch.rand(2, 4, 4),
        "b": torch.rand(2, 4, 4),
    }
    data["mask"][:, :, 2:] = 0
    gen = {k: torch.rand_like(v) for k, v in target.items()}
    expected_gen = {k: v.clone() for k, v in gen.items()}
    expected_gen["a"][:, :, 2:] = target["a"][:, :, 2:]
    assert not torch.allclose(gen["a"], expected_gen["a"])
    prescribed_gen = prescriber(data, gen, target)
    for name in gen:
        torch.testing.assert_allclose(prescribed_gen[name], expected_gen[name])
    # non-integer valued mask
    prescriber = Prescriber(prescribed_name="a", mask_name="mask", mask_value=1)
    data["mask"] = torch.zeros(2, 4, 4, dtype=torch.float32) + 0.1
    data["mask"][:, :, 2:] = 0.7
    prescribed_gen = prescriber(data, gen, target)
    for name in gen:
        torch.testing.assert_allclose(prescribed_gen[name], expected_gen[name])
//  [python]
def test_prescriber():
    """Test that the prescriber overwrites the generated data in the masked region."""
    prescriber = Prescriber(prescribed_name="a", mask_name="mask", mask_value=0)
    data = {
        "a": torch.rand(2, 4, 4),
        "b": torch.rand(2, 4, 4),
        "mask": torch.ones(2, 4, 4),
    }
    target = {
        "a": torch.rand(2, 4, 4),
        "b": torch.rand(2, 4, 4),
    }
    data["mask"][:, :, 2:] = 0
    gen = {k: torch.rand_like(v) for k, v in target.items()}
    expected_gen = {k: v.clone() for k, v in gen.items()}
    expected_gen["a"][:, :, 2:] = target["a"][:, :, 2:]
    assert not torch.allclose(gen["a"], expected_gen["a"])
    prescribed_gen = prescriber(data, gen, target)
    for name in gen:
        torch.testing.assert_close(prescribed_gen[name], expected_gen[name])
    # non-integer valued mask
    prescriber = Prescriber(prescribed_name="a", mask_name="mask", mask_value=1)
    data["mask"] = torch.zeros(2, 4, 4, dtype=torch.float32) + 0.1
    data["mask"][:, :, 2:] = 0.7
    prescribed_gen = prescriber(data, gen, target)
    for name in gen:
        torch.testing.assert_close(prescribed_gen[name], expected_gen[name])
//  [python]
def test_get_tensor_mean_n_stds_min_max():
    pytest_raises_wrapper(ValueError, 'Expecting ValueError with n_stds = 0',
                          qu.get_tensor_mean_n_stds_min_max, test_tensor, n_stds=0)

    mean = torch.tensor(-16.)
    std = torch.tensor(62.87447738647461)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor)
    torch.testing.assert_allclose(t_min, mean - std)
    torch.testing.assert_allclose(t_max, mean + std)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor, n_stds=2)
    torch.testing.assert_allclose(t_min, torch.tensor(-95.))
    torch.testing.assert_allclose(t_max, torch.tensor(87.))
//  [python]
def test_get_tensor_mean_n_stds_min_max():
    pytest_raises_wrapper(ValueError, 'Expecting ValueError with n_stds = 0',
                          qu.get_tensor_mean_n_stds_min_max, test_tensor, n_stds=0)

    mean = torch.tensor(-16.)
    std = torch.tensor(62.87447738647461)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor)
    torch.testing.assert_close(t_min, mean - std)
    torch.testing.assert_close(t_max, mean + std)

    t_min, t_max = qu.get_tensor_mean_n_stds_min_max(test_tensor, n_stds=2)
    torch.testing.assert_close(t_min, torch.tensor(-95.))
    torch.testing.assert_close(t_max, torch.tensor(87.))
//  [python]
def test_skim_streaming():
    SEQ_LEN = 100
    num_spk = 2
    BS = 2
    separator = SkiMSeparator(
        input_dim=128,
        causal=True,
        num_spk=2,
        layer=2,
        unit=32,
        segment_size=100,
    )
    separator.eval()
    input_feature = torch.randn((BS, SEQ_LEN, 128))
    ilens = torch.LongTensor([SEQ_LEN] * BS)
    with torch.no_grad():
        seq_output, _, _ = separator.forward(input_feature, ilens=ilens)

        state = None
        stream_outputs = []
        for i in range(SEQ_LEN):
            frame = input_feature[:, i : i + 1, :]
            frame_out, state, _ = separator.forward_streaming(frame, state)
            stream_outputs.append(frame_out)
        for i in range(SEQ_LEN):
            for s in range(num_spk):
                torch.testing.assert_allclose(
                    stream_outputs[i][s], seq_output[s][:, i : i + 1, :]
                )
//  [python]
def test_skim_streaming():
    SEQ_LEN = 100
    num_spk = 2
    BS = 2
    separator = SkiMSeparator(
        input_dim=128,
        causal=True,
        num_spk=2,
        layer=2,
        unit=32,
        segment_size=100,
    )
    separator.eval()
    input_feature = torch.randn((BS, SEQ_LEN, 128))
    ilens = torch.LongTensor([SEQ_LEN] * BS)
    with torch.no_grad():
        seq_output, _, _ = separator.forward(input_feature, ilens=ilens)

        state = None
        stream_outputs = []
        for i in range(SEQ_LEN):
            frame = input_feature[:, i : i + 1, :]
            frame_out, state, _ = separator.forward_streaming(frame, state)
            stream_outputs.append(frame_out)
        for i in range(SEQ_LEN):
            for s in range(num_spk):
                torch.testing.assert_close(
                    stream_outputs[i][s], seq_output[s][:, i : i + 1, :]
                )
//  [python]
def test_set_value(self, dim_reward):
        value_function = TabularQFunction(
            num_states=4, num_actions=2, dim_reward=(dim_reward,)
        )
        value_function.set_value(2, 1, torch.tensor([1.0]))
        table = torch.zeros(2, 4, dim_reward)
        table[1, 2, :] = 1.0 * torch.ones(dim_reward)
        torch.testing.assert_allclose(value_function.table, table)
//  [python]
def test_set_value(self, dim_reward):
        value_function = TabularQFunction(
            num_states=4, num_actions=2, dim_reward=(dim_reward,)
        )
        value_function.set_value(2, 1, torch.tensor([1.0]))
        table = torch.zeros(2, 4, dim_reward)
        table[1, 2, :] = 1.0 * torch.ones(dim_reward)
        torch.testing.assert_close(value_function.table, table)
//  [python]
def test(xlist, blist, rtol=1e-3, atol=1e-5, tolerate_small_mismatch=True):
    for a, b in zip(xlist, blist):
        try:
            torch.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
        except AssertionError as error:
            if tolerate_small_mismatch:
                print(error)
            else:
                raise
//  [python]
def test(xlist, blist, rtol=1e-3, atol=1e-5, tolerate_small_mismatch=True):
    for a, b in zip(xlist, blist):
        try:
            torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
        except AssertionError as error:
            if tolerate_small_mismatch:
                print(error)
            else:
                raise
//  [python]
def test_groupby(stable: bool):
    n, k, d = 1000, 3, 10

    x = torch.randn(n, d)
    z = torch.randint(0, k, (n,))
    grouped = groupby(x, z, stable=stable)

    # Check that we can coalesce the groups back together
    torch.testing.assert_allclose(x, grouped.coalesce())
    torch.testing.assert_allclose(
        x + 1,
        grouped.map(lambda _, g: g + 1).coalesce(),
    )

    # We only expect these to be the same when stable=True
    if stable:
        for label in torch.unique(z):
            torch.testing.assert_allclose(
                x[z == label],
                grouped.groups[label],
            )
//  [python]
def test_groupby(stable: bool):
    n, k, d = 1000, 3, 10

    x = torch.randn(n, d)
    z = torch.randint(0, k, (n,))
    grouped = groupby(x, z, stable=stable)

    # Check that we can coalesce the groups back together
    torch.testing.assert_close(x, grouped.coalesce())
    torch.testing.assert_close(
        x + 1,
        grouped.map(lambda _, g: g + 1).coalesce(),
    )

    # We only expect these to be the same when stable=True
    if stable:
        for label in torch.unique(z):
            torch.testing.assert_close(
                x[z == label],
                grouped.groups[label],
            )
//  [python]
def test_weights_content_single():
    h = torch.tensor([10, 0, 10])[:, None, None].float()

    attn = Weights(3, 1, position=False, norm=False)
    attn.k.weight[:] = +1.
    attn.q.weight[:] = +1.
    attn.k_bias[:] = 0.
    what = attn(h)

    w = torch.tensor([[0, 0, 0], [1, 0, 0], [1, 0, 0]]).float()
    torch.testing.assert_allclose(w, what[:, :, 0, 0])
//  [python]
def test_weights_content_single():
    h = torch.tensor([10, 0, 10])[:, None, None].float()

    attn = Weights(3, 1, position=False, norm=False)
    attn.k.weight[:] = +1.
    attn.q.weight[:] = +1.
    attn.k_bias[:] = 0.
    what = attn(h)

    w = torch.tensor([[0, 0, 0], [1, 0, 0], [1, 0, 0]]).float()
    torch.testing.assert_close(w, what[:, :, 0, 0])
//  [python]
def test_weights_content_multi():
    h0 = torch.tensor([10, 0])[:, None, None].float()
    h1 = torch.tensor([10,])[:, None, None].float()

    attn = Weights(3, 1, position=False, norm=False)
    attn.k.weight[:] = +1.
    attn.q.weight[:] = +1.
    attn.k_bias[:] = 0.

    attn(h0)
    what = attn(h1)

    w = torch.tensor([[1, 0, 0]]).float()
    torch.testing.assert_allclose(w, what[:, :, 0, 0])
//  [python]
def test_weights_content_multi():
    h0 = torch.tensor([10, 0])[:, None, None].float()
    h1 = torch.tensor([10,])[:, None, None].float()

    attn = Weights(3, 1, position=False, norm=False)
    attn.k.weight[:] = +1.
    attn.q.weight[:] = +1.
    attn.k_bias[:] = 0.

    attn(h0)
    what = attn(h1)

    w = torch.tensor([[1, 0, 0]]).float()
    torch.testing.assert_close(w, what[:, :, 0, 0])
//  [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)
        )  # noqa: C402
        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)
        )  # noqa: C402
        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 testUnrollEqualityBackward(self):
        """Check if metagrads match for target/ref net."""
        for test_it in range(5):
            with higher.innerloop_ctx(self.target_net,
                                      self.opt) as (fnet, diffopt):
                ref_out, target_out = self._joint_inner_loop(
                    fnet, diffopt=diffopt, num_steps=10
                )
                ref_meta_loss = ref_out[0]
                target_meta_loss = target_out[0]

                ref_metagrads = torch.autograd.grad(
                    ref_meta_loss, self.reference_net.parameters()
                )
                target_metagrads = torch.autograd.grad(
                    target_meta_loss, fnet.parameters(time=0)
                )

                # Check that metagrads align
                for rg, tg in zip(ref_metagrads, target_metagrads):
                    torch.testing.assert_allclose(rg, tg)