jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_blur_equivalence(kernel_size, sigma, channels):
    for _ in range(10):
        input_tensor = torch.randn((3, channels, 128, 128))
        kornia = korniaGaussianBlur2d(kernel_size, sigma, separable=False)
        blur_kornia = kornia(input_tensor)
        gaussian = GaussianBlur2d(kernel_size=kernel_size, sigma=sigma, channels=channels)
        blur_gaussian = gaussian(input_tensor)
        torch.testing.assert_close(blur_kornia, blur_gaussian)
//  [python]
def toy_problem(seed=0, n_points=500):
    torch.random.manual_seed(seed)
    aa = torch.randn(3) / 10
    t = torch.randn(3) / 5
    T_w2q = Pose.from_aa(aa, t)

    w, h = 640, 480
    fx, fy = 300., 350.
    cx, cy = w/2, h/2
    radial = [0.1, 0.01]
    camera = Camera(torch.tensor([w, h, fx, fy, cx, cy] + radial)).float()
    torch.testing.assert_allclose((w, h), camera.size.long())
    torch.testing.assert_allclose((fx, fy), camera.f)
    torch.testing.assert_allclose((cx, cy), camera.c)

    p3D = torch.randn(n_points, 3)
    p3D[:, -1] += 2

    dim = 16
    F_ref = torch.randn(n_points, dim)
    F_query = torch.randn(dim, h, w)

    return T_w2q, camera, p3D, F_ref, F_query
//  [python]
def toy_problem(seed=0, n_points=500):
    torch.random.manual_seed(seed)
    aa = torch.randn(3) / 10
    t = torch.randn(3) / 5
    T_w2q = Pose.from_aa(aa, t)

    w, h = 640, 480
    fx, fy = 300., 350.
    cx, cy = w/2, h/2
    radial = [0.1, 0.01]
    camera = Camera(torch.tensor([w, h, fx, fy, cx, cy] + radial)).float()
    torch.testing.assert_close((w, h), camera.size.long())
    torch.testing.assert_close((fx, fy), camera.f)
    torch.testing.assert_close((cx, cy), camera.c)

    p3D = torch.randn(n_points, 3)
    p3D[:, -1] += 2

    dim = 16
    F_ref = torch.randn(n_points, dim)
    F_query = torch.randn(dim, h, w)

    return T_w2q, camera, p3D, F_ref, F_query
//  [python]
def test_maskedconvk5(self):
        torch.testing.assert_allclose(self.convk5(self.data),
                                      torch.tensor([[17.,  -16.,   32.,  -13.,   65.],
                                                    [-15.,   78.,  -23.,  107.,   12.],
                                                    [-20.,   28.,  -29.,   49.,  -18.],
                                                    [85.,  -20.,  100.,   -1.,  -18.],
                                                    [42.,  -15.,   55.,  -17.,   85.]],
                                                   dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def test_maskedconvk5(self):
        torch.testing.assert_close(self.convk5(self.data),
                                      torch.tensor([[17.,  -16.,   32.,  -13.,   65.],
                                                    [-15.,   78.,  -23.,  107.,   12.],
                                                    [-20.,   28.,  -29.,   49.,  -18.],
                                                    [85.,  -20.,  100.,   -1.,  -18.],
                                                    [42.,  -15.,   55.,  -17.,   85.]],
                                                   dtype=torch.float).unsqueeze(0).unsqueeze(0))
//  [python]
def assert_weights_same(torch_module, tf_module, atol=1e-6, rtol=1e-5):
    for th_weight, tf_variable in zip_th_tf_parameters(torch_module, tf_module, permute_torch=True):
        tf_weight = torch.tensor(tf_variable.numpy())
        torch.testing.assert_allclose(tf_weight, th_weight, atol=atol, rtol=rtol)
//  [python]
def assert_weights_same(torch_module, tf_module, atol=1e-6, rtol=1e-5):
    for th_weight, tf_variable in zip_th_tf_parameters(torch_module, tf_module, permute_torch=True):
        tf_weight = torch.tensor(tf_variable.numpy())
        torch.testing.assert_close(tf_weight, th_weight, atol=atol, rtol=rtol)
//  [python]
def test_dmi_bundle(self):

        bundle = DependencyAwareMutualInformationBundle()
        mig = MutualInformationGap(fill_reg_dim=True)
        dmig = DependencyAwareMutualInformationGap()
        dlig = DependencyAwareLatentInformationGap()
        xmig = DependencyBlindMutualInformationGap()

        z = torch.randn(16, 16)
        a = torch.randn(16, 3)

        bundle.update(z=z, a=a)
        mig.update(z, a)
        dmig.update(z, a)
        dlig.update(z, a)
        xmig.update(z, a)

        bundle_out = bundle.compute()
        indiv_out = {
            "MIG": mig.compute(),
            "DMIG": dmig.compute(),
            "DLIG": dlig.compute(),
            "XMIG": xmig.compute(),
        }

        for key in ["MIG", "DMIG", "DLIG", "XMIG"]:
            torch.testing.assert_allclose(bundle_out[key], indiv_out[key])
//  [python]
def test_dmi_bundle(self):

        bundle = DependencyAwareMutualInformationBundle()
        mig = MutualInformationGap(fill_reg_dim=True)
        dmig = DependencyAwareMutualInformationGap()
        dlig = DependencyAwareLatentInformationGap()
        xmig = DependencyBlindMutualInformationGap()

        z = torch.randn(16, 16)
        a = torch.randn(16, 3)

        bundle.update(z=z, a=a)
        mig.update(z, a)
        dmig.update(z, a)
        dlig.update(z, a)
        xmig.update(z, a)

        bundle_out = bundle.compute()
        indiv_out = {
            "MIG": mig.compute(),
            "DMIG": dmig.compute(),
            "DLIG": dlig.compute(),
            "XMIG": xmig.compute(),
        }

        for key in ["MIG", "DMIG", "DLIG", "XMIG"]:
            torch.testing.assert_close(bundle_out[key], indiv_out[key])
//  [python]
def test_masked_fill(self):
        dtypes = [
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        sizes = [(2,), (4, 4)]
        for self_dtype, device, scalar_val, size in product(dtypes, self.devices, [0.4, 3], sizes):
            input_v = self.data_for(self_dtype, device, size=size)
            mask = self.data_for(torch.bool, device, size=size)

            def fn(input_v, mask):
                return torch.masked_fill(input_v, mask, scalar_val)
            ref = fn(input_v, mask)
            try:
                t = torch.jit.trace(fn, (input_v, mask))
                torch.testing.assert_allclose(ref, t(input_v, mask))
                print(torch.jit.last_executed_optimized_graph())
                self.assertLastGraphAllFused()
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(self_dtype), op.__name__, device, str(size)])
                )
//  [python]
def test_masked_fill(self):
        dtypes = [
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        sizes = [(2,), (4, 4)]
        for self_dtype, device, scalar_val, size in product(dtypes, self.devices, [0.4, 3], sizes):
            input_v = self.data_for(self_dtype, device, size=size)
            mask = self.data_for(torch.bool, device, size=size)

            def fn(input_v, mask):
                return torch.masked_fill(input_v, mask, scalar_val)
            ref = fn(input_v, mask)
            try:
                t = torch.jit.trace(fn, (input_v, mask))
                torch.testing.assert_close(ref, t(input_v, mask))
                print(torch.jit.last_executed_optimized_graph())
                self.assertLastGraphAllFused()
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(self_dtype), op.__name__, device, str(size)])
                )
//  [python]
def test_dataloaders(self):
        dataset = ForwardShapenetDataset(self.config)
        dataset.create_dataloaders(MockModel(DictConfig({"conv_type": "DENSE"})), 2, False, 1, False)
        forward_set = dataset.test_dataloaders[0]
        for b in forward_set:
            self.assertEqual(b.origin_id.shape, (2, 2))

        sparseconfig = DictConfig({"dataroot": self.datadir, "category": "Airplane", "forward_category": "Airplane"})
        dataset = ForwardShapenetDataset(sparseconfig)
        dataset.create_dataloaders(MockModel(DictConfig({"conv_type": "PARTIAL_DENSE"})), 2, False, 1, False)
        forward_set = dataset.test_dataloaders[0]
        for b in forward_set:
            torch.testing.assert_allclose(b.origin_id, torch.tensor([0, 1, 2, 0, 1, 2, 3]))
            torch.testing.assert_allclose(b.sampleid, torch.tensor([0, 1]))
//  [python]
def test_dataloaders(self):
        dataset = ForwardShapenetDataset(self.config)
        dataset.create_dataloaders(MockModel(DictConfig({"conv_type": "DENSE"})), 2, False, 1, False)
        forward_set = dataset.test_dataloaders[0]
        for b in forward_set:
            self.assertEqual(b.origin_id.shape, (2, 2))

        sparseconfig = DictConfig({"dataroot": self.datadir, "category": "Airplane", "forward_category": "Airplane"})
        dataset = ForwardShapenetDataset(sparseconfig)
        dataset.create_dataloaders(MockModel(DictConfig({"conv_type": "PARTIAL_DENSE"})), 2, False, 1, False)
        forward_set = dataset.test_dataloaders[0]
        for b in forward_set:
            torch.testing.assert_close(b.origin_id, torch.tensor([0, 1, 2, 0, 1, 2, 3]))
            torch.testing.assert_close(b.sampleid, torch.tensor([0, 1]))
//  [python]
def test_cupy_knn_points_results():
    d2_0, idx_0, nn_0 = knn_points(pts, pts, K=K, return_nn=False, return_sorted=True)
    d2_1, idx_1, nn_1 = raw_cupy_knn_points(pts, pts, K=K, return_nn=False, return_sorted=True)

    torch.testing.assert_allclose(d2_0, d2_1)
    torch.testing.assert_allclose(idx_0, idx_1)
//  [python]
def test_cupy_knn_points_results():
    d2_0, idx_0, nn_0 = knn_points(pts, pts, K=K, return_nn=False, return_sorted=True)
    d2_1, idx_1, nn_1 = raw_cupy_knn_points(pts, pts, K=K, return_nn=False, return_sorted=True)

    torch.testing.assert_close(d2_0, d2_1)
    torch.testing.assert_close(idx_0, idx_1)
//  [python]
def test_unary_ops(self):
        def apply(fn):
            return lambda x: fn(x)

        unary_ops = [
            torch.lgamma,
            torch.sigmoid,
            torch.reciprocal,
            torch.neg,
            torch.relu,
            torch.log,
            torch.log10,
            torch.log1p,
            torch.log2,
            torch.exp,
            torch.expm1,
            torch.erf,
            torch.erfc,
            torch.cos,
            torch.sin,
            torch.tan,
            torch.acos,
            torch.asin,
            torch.cosh,
            torch.sinh,
            torch.atan,
            torch.tanh,
            torch.sqrt,
            torch.rsqrt,
            torch.abs,
            torch.ceil,
            torch.floor,
            torch.round,
            torch.trunc,
            torch.frac,
            lambda x: torch.threshold(x, 0, -10),
            lambda x: torch.clamp(x, -10, 10),
        ]
        sizes = [(1,), (2,), (4, 4)]
        for dtype, op, device, size in product(self.dtypes, unary_ops, self.devices, sizes):
            try:
                x = self.data_for(dtype, device, size=size)
                fn = apply(op)
                ref = fn(x)
            except Exception:
                # If eager mode doesn't support a dtype/op/device combo,
                # neither does the fuser.  Catch everything to avoid needing to
                # guess what errors might be thrown by eager.
                continue
            try:
                t = torch.jit.trace(fn, (x,))
                torch.testing.assert_allclose(ref, t(x))
                self.assertAllFused(t.graph_for(x))
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(dtype), op.__name__, device, str(size)])
                )
//  [python]
def test_unary_ops(self):
        def apply(fn):
            return lambda x: fn(x)

        unary_ops = [
            torch.lgamma,
            torch.sigmoid,
            torch.reciprocal,
            torch.neg,
            torch.relu,
            torch.log,
            torch.log10,
            torch.log1p,
            torch.log2,
            torch.exp,
            torch.expm1,
            torch.erf,
            torch.erfc,
            torch.cos,
            torch.sin,
            torch.tan,
            torch.acos,
            torch.asin,
            torch.cosh,
            torch.sinh,
            torch.atan,
            torch.tanh,
            torch.sqrt,
            torch.rsqrt,
            torch.abs,
            torch.ceil,
            torch.floor,
            torch.round,
            torch.trunc,
            torch.frac,
            lambda x: torch.threshold(x, 0, -10),
            lambda x: torch.clamp(x, -10, 10),
        ]
        sizes = [(1,), (2,), (4, 4)]
        for dtype, op, device, size in product(self.dtypes, unary_ops, self.devices, sizes):
            try:
                x = self.data_for(dtype, device, size=size)
                fn = apply(op)
                ref = fn(x)
            except Exception:
                # If eager mode doesn't support a dtype/op/device combo,
                # neither does the fuser.  Catch everything to avoid needing to
                # guess what errors might be thrown by eager.
                continue
            try:
                t = torch.jit.trace(fn, (x,))
                torch.testing.assert_close(ref, t(x))
                self.assertAllFused(t.graph_for(x))
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(dtype), op.__name__, device, str(size)])
                )
//  [python]
def test_cupy_knn_results():
    d2, idx, nn = knn_points(pts, pts, K=K, return_nn=False, return_sorted=True)

    pts_cupy = cupy.asarray(pts)
    lbvh = cupy_knn.LBVHIndex(leaf_size=32,
                              compact=False,
                              shrink_to_fit=False,
                              sort_queries=True)
    lbvh.build(pts_cupy)
    lbvh.prepare_knn_default(K, radius=None)
    idx_cupy, d2_cupy, nn_cupy = lbvh.query_knn(pts_cupy)
    idx_tensor = torch.as_tensor(idx_cupy.astype(cupy.int64), device=device)[None]  # !: BATCH
    d2_tensor = torch.as_tensor(d2_cupy, device=device)[None]
    nn_tensor = torch.as_tensor(nn_cupy.astype(cupy.int64), device=device)[None]

    torch.testing.assert_allclose(d2, d2_tensor)
    torch.testing.assert_allclose(idx, idx_tensor)
//  [python]
def test_cupy_knn_results():
    d2, idx, nn = knn_points(pts, pts, K=K, return_nn=False, return_sorted=True)

    pts_cupy = cupy.asarray(pts)
    lbvh = cupy_knn.LBVHIndex(leaf_size=32,
                              compact=False,
                              shrink_to_fit=False,
                              sort_queries=True)
    lbvh.build(pts_cupy)
    lbvh.prepare_knn_default(K, radius=None)
    idx_cupy, d2_cupy, nn_cupy = lbvh.query_knn(pts_cupy)
    idx_tensor = torch.as_tensor(idx_cupy.astype(cupy.int64), device=device)[None]  # !: BATCH
    d2_tensor = torch.as_tensor(d2_cupy, device=device)[None]
    nn_tensor = torch.as_tensor(nn_cupy.astype(cupy.int64), device=device)[None]

    torch.testing.assert_close(d2, d2_tensor)
    torch.testing.assert_close(idx, idx_tensor)
//  [python]
def test_subsample_with_padding(self):
        encoder = S2TTransformerEncoder(self.args, self.src_dict)
        encoder.eval()
        torch.manual_seed(0)
        fake_sample = torch.rand(2, 40, 4)
        fake_sample[1, 20:, :] = 0
        fake_lengths = torch.LongTensor([40, 20])
        fake_sample2 = fake_sample[1:, :20, :]
        subsample_out = encoder.subsample(fake_sample, fake_lengths)
        subsample_out2 = encoder.subsample(fake_sample2, fake_lengths[1:])
        # Sum of padded elements must be equal to 0
        self.assertTrue((subsample_out[0][5:, 1, :].sum() == 0).item())
        # Non padded elements of the processed batch must be equal to the unpadded sample
        torch.testing.assert_allclose(
            subsample_out[0][:5, 1, :],
            subsample_out2[0][:, 0, :])
        encoder_out = encoder(fake_sample, fake_lengths)
        encoder_out2 = encoder(fake_sample2, fake_lengths[1:])
        # Ensure encoder output
        torch.testing.assert_allclose(
            encoder_out['encoder_out'][0][:5, 1, :],
            encoder_out2['encoder_out'][0][:, 0, :])
//  [python]
def test_subsample_with_padding(self):
        encoder = S2TTransformerEncoder(self.args, self.src_dict)
        encoder.eval()
        torch.manual_seed(0)
        fake_sample = torch.rand(2, 40, 4)
        fake_sample[1, 20:, :] = 0
        fake_lengths = torch.LongTensor([40, 20])
        fake_sample2 = fake_sample[1:, :20, :]
        subsample_out = encoder.subsample(fake_sample, fake_lengths)
        subsample_out2 = encoder.subsample(fake_sample2, fake_lengths[1:])
        # Sum of padded elements must be equal to 0
        self.assertTrue((subsample_out[0][5:, 1, :].sum() == 0).item())
        # Non padded elements of the processed batch must be equal to the unpadded sample
        torch.testing.assert_close(
            subsample_out[0][:5, 1, :],
            subsample_out2[0][:, 0, :])
        encoder_out = encoder(fake_sample, fake_lengths)
        encoder_out2 = encoder(fake_sample2, fake_lengths[1:])
        # Ensure encoder output
        torch.testing.assert_close(
            encoder_out['encoder_out'][0][:5, 1, :],
            encoder_out2['encoder_out'][0][:, 0, :])
//  [python]
def test_world_to_car(self) -> None:
        batch = dummy_batch()
        out = batch.world_to_car(1)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_allclose(out, batch.cam_T[:, 1])
//  [python]
def test_world_to_car(self) -> None:
        batch = dummy_batch()
        out = batch.world_to_car(1)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_close(out, batch.cam_T[:, 1])
//  [python]
relevance.register_function(torch.nn.Linear, relevance.LinearEpsilon)

graphs = tg.GraphBatch.from_graphs([tg.Graph(
    node_features=torch.rand(10, 9),
    edge_features=torch.rand(6, 5),
    senders=torch.tensor([0, 0, 1, 2, 4, 5]),
    receivers=torch.tensor([1, 2, 2, 4, 3, 3]),
    global_features=torch.rand(7)
)])

print('bias ef sf rf gf out input'.replace(' ', '\t'))
print('-' * 45)

for bias, ef, sf, rf, gf in itertools.product(
        [True, False, 'zero'],
        [5, None],
        [9, None],
        [9, None],
        [7, None]
):
    if ef is sf is rf is gf is None:
        continue
    net = tg.EdgeLinear(3, edge_features=ef, sender_features=sf, receiver_features=rf, global_features=gf, bias=bool(bias))
    if bias == 'zero':
        net.bias.zero_()

    ctx = {}
    out = relevance.forward_relevance(net, graphs, ctx=ctx)
    torch.testing.assert_allclose(out.edge_features, net(graphs).edge_features)

    rel_out = graphs.evolve(
        edge_features=torch.ones_like(out.edge_features) * (out.edge_features != 0).float()
    )

    rel_in = relevance.backward_relevance(net, rel_out, ctx=ctx)

    # If bias==0 then relevance is conserved at a graph level
    print(
        {True: '  x ', False: '  - ', 'zero': '  0 '}[bias], ef if ef else '-', sf if sf else '-',
        rf if rf else '-', gf if gf else '-',
        rel_out.edge_features.sum().item(),
        ((rel_in.global_features.sum() if rel_in.global_features is not None else 0) +
         (rel_in.edge_features.sum() if rel_in.edge_features is not None else 0) +
         (rel_in.node_features.sum() if rel_in.node_features is not None else 0)).item(),
        sep='\t'
    )
//  [python]
relevance.register_function(torch.nn.Linear, relevance.LinearEpsilon)

graphs = tg.GraphBatch.from_graphs([tg.Graph(
    node_features=torch.rand(10, 9),
    edge_features=torch.rand(6, 5),
    senders=torch.tensor([0, 0, 1, 2, 4, 5]),
    receivers=torch.tensor([1, 2, 2, 4, 3, 3]),
    global_features=torch.rand(7)
)])

print('bias ef sf rf gf out input'.replace(' ', '\t'))
print('-' * 45)

for bias, ef, sf, rf, gf in itertools.product(
        [True, False, 'zero'],
        [5, None],
        [9, None],
        [9, None],
        [7, None]
):
    if ef is sf is rf is gf is None:
        continue
    net = tg.EdgeLinear(3, edge_features=ef, sender_features=sf, receiver_features=rf, global_features=gf, bias=bool(bias))
    if bias == 'zero':
        net.bias.zero_()

    ctx = {}
    out = relevance.forward_relevance(net, graphs, ctx=ctx)
    torch.testing.assert_close(out.edge_features, net(graphs).edge_features)

    rel_out = graphs.evolve(
        edge_features=torch.ones_like(out.edge_features) * (out.edge_features != 0).float()
    )

    rel_in = relevance.backward_relevance(net, rel_out, ctx=ctx)

    # If bias==0 then relevance is conserved at a graph level
    print(
        {True: '  x ', False: '  - ', 'zero': '  0 '}[bias], ef if ef else '-', sf if sf else '-',
        rf if rf else '-', gf if gf else '-',
        rel_out.edge_features.sum().item(),
        ((rel_in.global_features.sum() if rel_in.global_features is not None else 0) +
         (rel_in.edge_features.sum() if rel_in.edge_features is not None else 0) +
         (rel_in.node_features.sum() if rel_in.node_features is not None else 0)).item(),
        sep='\t'
    )
//  [python]
def test_bert_embeddings():
    model_name = "bert_base_embeddings"
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    qconfig = [
        ("BACKEND", "tensorrt"),
        ("W.QSCHEME", "per-channel-symmetric"),
        ("W.QUANTIZER.TYPE", "uniform"),
        ("W.QUANTIZER.BIT", 8),
        ("W.OBSERVER.TYPE", "MINMAX"),
        ("A.QSCHEME", "per-tensor-symmetric"),
        ("A.QUANTIZER.TYPE", "uniform"),
        ("A.QUANTIZER.BIT", 8),
        ("A.OBSERVER.TYPE", "MINMAX"),
    ]
    qconfig = [j for i in qconfig for j in i]
    qconfig = build_qconfig(qconfig)

    model_config = BertConfig()
    model = BertEmbeddings(model_config)
    model.seq_length = model_config.seq_length

    qmodel = QuantModel(model, qconfig)

    input_ids = torch.zeros(1, 32)
    input_ids[:, :10] = torch.tensor(
        [
            101,
            3533,
            2003,
            12283,
            2084,
            1045,
            2228,
            2984,
            1012,
            102,
        ]
    )
    input_ids = input_ids.long()
    token_type_ids = torch.zeros_like(input_ids, dtype=torch.long)

    model.eval()
    qmodel.eval()
    out1 = model(input_ids, token_type_ids)
    out2 = qmodel(input_ids, token_type_ids)
    torch.testing.assert_allclose(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test_bert_embeddings():
    model_name = "bert_base_embeddings"
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    qconfig = [
        ("BACKEND", "tensorrt"),
        ("W.QSCHEME", "per-channel-symmetric"),
        ("W.QUANTIZER.TYPE", "uniform"),
        ("W.QUANTIZER.BIT", 8),
        ("W.OBSERVER.TYPE", "MINMAX"),
        ("A.QSCHEME", "per-tensor-symmetric"),
        ("A.QUANTIZER.TYPE", "uniform"),
        ("A.QUANTIZER.BIT", 8),
        ("A.OBSERVER.TYPE", "MINMAX"),
    ]
    qconfig = [j for i in qconfig for j in i]
    qconfig = build_qconfig(qconfig)

    model_config = BertConfig()
    model = BertEmbeddings(model_config)
    model.seq_length = model_config.seq_length

    qmodel = QuantModel(model, qconfig)

    input_ids = torch.zeros(1, 32)
    input_ids[:, :10] = torch.tensor(
        [
            101,
            3533,
            2003,
            12283,
            2084,
            1045,
            2228,
            2984,
            1012,
            102,
        ]
    )
    input_ids = input_ids.long()
    token_type_ids = torch.zeros_like(input_ids, dtype=torch.long)

    model.eval()
    qmodel.eval()
    out1 = model(input_ids, token_type_ids)
    out2 = qmodel(input_ids, token_type_ids)
    torch.testing.assert_close(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test_car_to_world(self) -> None:
        batch = dummy_batch()
        out = batch.car_to_world(1)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_allclose(out, batch.cam_T[:, 1].pinverse())
//  [python]
def test_car_to_world(self) -> None:
        batch = dummy_batch()
        out = batch.car_to_world(1)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_close(out, batch.cam_T[:, 1].pinverse())
//  [python]
def test_world_to_cam(self) -> None:
        batch = dummy_batch()
        cam = "left"
        frame = 1
        target = batch.T[cam].pinverse().matmul(batch.cam_T[:, frame])
        out = batch.world_to_cam(cam, frame)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_allclose(out, target)
//  [python]
def test_world_to_cam(self) -> None:
        batch = dummy_batch()
        cam = "left"
        frame = 1
        target = batch.T[cam].pinverse().matmul(batch.cam_T[:, frame])
        out = batch.world_to_cam(cam, frame)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_close(out, target)
//  [python]
def test_cam_to_world(self) -> None:
        torch.manual_seed(0)
        torch.use_deterministic_algorithms(True, warn_only=True)

        batch = dummy_batch()
        cam = "left"
        frame = 1
        target = batch.T[cam].inverse().matmul(batch.cam_T[:, frame]).inverse()
        out = batch.cam_to_world(cam, frame)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_allclose(out, target)
//  [python]
def test_cam_to_world(self) -> None:
        torch.manual_seed(0)
        torch.use_deterministic_algorithms(True, warn_only=True)

        batch = dummy_batch()
        cam = "left"
        frame = 1
        target = batch.T[cam].inverse().matmul(batch.cam_T[:, frame]).inverse()
        out = batch.cam_to_world(cam, frame)
        self.assertEqual(out.shape, (2, 4, 4))
        torch.testing.assert_close(out, target)
//  [python]
def test_grid_image(self) -> None:
        batch = dummy_batch()
        img = batch.grid_image("left", 1)
        self.assertEqual(img.data.shape, (2, 3, 48, 64))
        mask = img.mask
        self.assertIsNotNone(mask)
        self.assertEqual(mask.shape, (2, 1, 48, 64))
        torch.testing.assert_allclose(img.time, batch.frame_time[:, 1])
        self.assertFalse(img.camera.in_ndc())
//  [python]
def test_grid_image(self) -> None:
        batch = dummy_batch()
        img = batch.grid_image("left", 1)
        self.assertEqual(img.data.shape, (2, 3, 48, 64))
        mask = img.mask
        self.assertIsNotNone(mask)
        self.assertEqual(mask.shape, (2, 1, 48, 64))
        torch.testing.assert_close(img.time, batch.frame_time[:, 1])
        self.assertFalse(img.camera.in_ndc())
//  [python]
def test_output(self, means):
        mean1, mean2, mean = means

        test_x = torch.linspace(0, 2, 5).unsqueeze(-1)
        test_x = test_x.expand(2, *test_x.shape)

        res = mean(test_x)

        torch.testing.assert_allclose(res[0], mean1(test_x[0]))
        torch.testing.assert_allclose(res[1], mean2(test_x[1]))

        torch.testing.assert_allclose(res[0], 1)
        torch.testing.assert_allclose(res[1], 2)
//  [python]
def test_output(self, means):
        mean1, mean2, mean = means

        test_x = torch.linspace(0, 2, 5).unsqueeze(-1)
        test_x = test_x.expand(2, *test_x.shape)

        res = mean(test_x)

        torch.testing.assert_close(res[0], mean1(test_x[0]))
        torch.testing.assert_close(res[1], mean2(test_x[1]))

        torch.testing.assert_close(res[0], 1)
        torch.testing.assert_close(res[1], 2)
//  [python]
def test_1d(self):
        x = torch.randn((10, 3))
        A = torch.randn((1, 3))

        mean = LinearMean(A, trainable=False)
        # Make sure matrix is not trainable
        assert not list(mean.parameters())
        assert mean.batch_size == 1

        out = mean(x)

        torch.testing.assert_allclose(out, (x @ A.t()).t()[0])
//  [python]
def test_1d(self):
        x = torch.randn((10, 3))
        A = torch.randn((1, 3))

        mean = LinearMean(A, trainable=False)
        # Make sure matrix is not trainable
        assert not list(mean.parameters())
        assert mean.batch_size == 1

        out = mean(x)

        torch.testing.assert_close(out, (x @ A.t()).t()[0])
//  [python]
def test_output(self, covariances):
        cov1, cov2, cov = covariances

        test_x = torch.linspace(0, 2, 5).unsqueeze(-1)
        test_x = test_x.expand(2, *test_x.shape)

        res = cov(test_x).evaluate()

        torch.testing.assert_allclose(res[0], cov1(test_x[0]).evaluate())
        torch.testing.assert_allclose(res[1], cov2(test_x[1]).evaluate())
//  [python]
def test_output(self, covariances):
        cov1, cov2, cov = covariances

        test_x = torch.linspace(0, 2, 5).unsqueeze(-1)
        test_x = test_x.expand(2, *test_x.shape)

        res = cov(test_x).evaluate()

        torch.testing.assert_close(res[0], cov1(test_x[0]).evaluate())
        torch.testing.assert_close(res[1], cov2(test_x[1]).evaluate())
//  [python]
def test_log_probs_from_logits_and_actions(self):
        """Tests categorical_distribution."""
        batch_size = 2
        seq_len = 7
        action_dim = 3

        policy_logits = _shaped_arange(seq_len, batch_size, action_dim) + 10
        actions = torch.from_numpy(np.random.randint(0, action_dim - 1, size=(seq_len, batch_size))).long()

        categorical_distribution = distributions.categorical_distribution(policy_logits)
        action_log_probs_tensor = categorical_distribution.log_prob(actions)

        # Ground Truth
        # Using broadcasting to create a mask that indexes action logits
        action_index_mask = actions[..., None] == torch.arange(action_dim)

        def index_with_mask(array, mask):
            return array[mask].reshape(*array.shape[:-1])

        # Note: Normally log(softmax) is not a good idea because it's not
        # numerically stable. However, in this test we have well-behaved values.
        ground_truth_v = index_with_mask(np.log(_softmax(policy_logits)), action_index_mask)

        torch.testing.assert_allclose(ground_truth_v, action_log_probs_tensor)
//  [python]
def test_log_probs_from_logits_and_actions(self):
        """Tests categorical_distribution."""
        batch_size = 2
        seq_len = 7
        action_dim = 3

        policy_logits = _shaped_arange(seq_len, batch_size, action_dim) + 10
        actions = torch.from_numpy(np.random.randint(0, action_dim - 1, size=(seq_len, batch_size))).long()

        categorical_distribution = distributions.categorical_distribution(policy_logits)
        action_log_probs_tensor = categorical_distribution.log_prob(actions)

        # Ground Truth
        # Using broadcasting to create a mask that indexes action logits
        action_index_mask = actions[..., None] == torch.arange(action_dim)

        def index_with_mask(array, mask):
            return array[mask].reshape(*array.shape[:-1])

        # Note: Normally log(softmax) is not a good idea because it's not
        # numerically stable. However, in this test we have well-behaved values.
        ground_truth_v = index_with_mask(np.log(_softmax(policy_logits)), action_index_mask)

        torch.testing.assert_close(ground_truth_v, action_log_probs_tensor)
//  [python]
def test_multidim(self):
        x = torch.randn((2, 10, 3))
        A = torch.randn((2, 3))

        mean = LinearMean(A, trainable=False)
        out = mean(x)

        assert mean.batch_size == 2
        # A @ x.T = (x.T @ A.T).T  for each x. The latter also works for multiple x.
        torch.testing.assert_allclose(out[[0]], (x[0] @ A[[0]].t()).t())
        torch.testing.assert_allclose(out[[1]], (x[1] @ A[[1]].t()).t())
//  [python]
def test_multidim(self):
        x = torch.randn((2, 10, 3))
        A = torch.randn((2, 3))

        mean = LinearMean(A, trainable=False)
        out = mean(x)

        assert mean.batch_size == 2
        # A @ x.T = (x.T @ A.T).T  for each x. The latter also works for multiple x.
        torch.testing.assert_close(out[[0]], (x[0] @ A[[0]].t()).t())
        torch.testing.assert_close(out[[1]], (x[1] @ A[[1]].t()).t())
//  [python]
def test_single_output_gp(self):
        kernel = gpytorch.kernels.MaternKernel(nu=2.5, ard_num_dims=None, batch_size=1, active_dims=None, lengthscale_prior=None, param_transform=softplus, inv_param_transform=None, eps=1e-6)
        mean = LinearMean(torch.tensor([[0.5]]))
        likelihood = gpytorch.likelihoods.GaussianLikelihood()
        #likelihood.noise = torch.tensor(0.01 ** 2)

        train_x = torch.tensor([-0.5, -0.1, 0., 0.1, 1.])[:, None]
        train_y = 0.5 * train_x.t()

        model = MultiOutputGP(train_x, train_y, kernel, likelihood, mean=mean)
        model.eval()

        test_x = torch.linspace(-1, 2, 5)
        pred = model(test_x)

        true_mean = torch.tensor([-0.5, -0.125, 0.25, 0.6250, 1.0])[None, :]
        torch.testing.assert_allclose(pred.mean, true_mean)
//  [python]
def test_single_output_gp(self):
        kernel = gpytorch.kernels.MaternKernel(nu=2.5, ard_num_dims=None, batch_size=1, active_dims=None, lengthscale_prior=None, param_transform=softplus, inv_param_transform=None, eps=1e-6)
        mean = LinearMean(torch.tensor([[0.5]]))
        likelihood = gpytorch.likelihoods.GaussianLikelihood()
        #likelihood.noise = torch.tensor(0.01 ** 2)

        train_x = torch.tensor([-0.5, -0.1, 0., 0.1, 1.])[:, None]
        train_y = 0.5 * train_x.t()

        model = MultiOutputGP(train_x, train_y, kernel, likelihood, mean=mean)
        model.eval()

        test_x = torch.linspace(-1, 2, 5)
        pred = model(test_x)

        true_mean = torch.tensor([-0.5, -0.125, 0.25, 0.6250, 1.0])[None, :]
        torch.testing.assert_close(pred.mean, true_mean)
//  [python]
def test_rnn_streaming():
    SEQ_LEN = 100
    num_spk = 2
    BS = 2
    separator = RNNSeparator(input_dim=128, rnn_type="lstm", num_spk=num_spk)
    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_rnn_streaming():
    SEQ_LEN = 100
    num_spk = 2
    BS = 2
    separator = RNNSeparator(input_dim=128, rnn_type="lstm", num_spk=num_spk)
    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_collate_custom_array():
    EMBEDDING_SIZE = 300

    cuts = CutSet.from_json("test/fixtures/ljspeech/cuts.json")
    with TemporaryDirectory() as d, NumpyFilesWriter(d) as writer:
        expected_xvectors = []
        for cut in cuts:
            expected_xvectors.append(np.random.randn(EMBEDDING_SIZE).astype(np.float32))
            cut.xvector = writer.store_array(cut.id, expected_xvectors[-1])

        xvectors = collate_custom_field(cuts, "xvector")
        assert isinstance(xvectors, torch.Tensor)
        assert xvectors.dtype == torch.float32
        assert xvectors.shape == (len(cuts), EMBEDDING_SIZE)
        for idx, xvec in enumerate(expected_xvectors):
            torch.testing.assert_allclose(xvectors[idx], xvec)
//  [python]
def test_collate_custom_array():
    EMBEDDING_SIZE = 300

    cuts = CutSet.from_json("test/fixtures/ljspeech/cuts.json")
    with TemporaryDirectory() as d, NumpyFilesWriter(d) as writer:
        expected_xvectors = []
        for cut in cuts:
            expected_xvectors.append(np.random.randn(EMBEDDING_SIZE).astype(np.float32))
            cut.xvector = writer.store_array(cut.id, expected_xvectors[-1])

        xvectors = collate_custom_field(cuts, "xvector")
        assert isinstance(xvectors, torch.Tensor)
        assert xvectors.dtype == torch.float32
        assert xvectors.shape == (len(cuts), EMBEDDING_SIZE)
        for idx, xvec in enumerate(expected_xvectors):
            torch.testing.assert_close(xvectors[idx], xvec)
//  [python]
def example_dconv():
    input = torch.randn(2, 64, 128, 128).cuda()
    # wrap all things (offset and mask) in DCN
    dcn = DCN(64, 64, kernel_size=(3, 3), stride=1, padding=1, deformable_groups=2).cuda()
    traced_dcn = torch.jit.script(dcn, input)
    # print(dcn.weight.shape, input.shape)
    output = dcn(input)
    result = traced_dcn(input)

    targert = output.new(*output.size())
    targert.data.uniform_(-0.01, 0.01)
    error = (targert - output).mean()
    error.backward()
    print(output.shape)
    print(traced_dcn.code)
    torch.testing.assert_allclose(output, result)
//  [python]
def example_dconv():
    input = torch.randn(2, 64, 128, 128).cuda()
    # wrap all things (offset and mask) in DCN
    dcn = DCN(64, 64, kernel_size=(3, 3), stride=1, padding=1, deformable_groups=2).cuda()
    traced_dcn = torch.jit.script(dcn, input)
    # print(dcn.weight.shape, input.shape)
    output = dcn(input)
    result = traced_dcn(input)

    targert = output.new(*output.size())
    targert.data.uniform_(-0.01, 0.01)
    error = (targert - output).mean()
    error.backward()
    print(output.shape)
    print(traced_dcn.code)
    torch.testing.assert_close(output, result)
//  [python]
def check_zero_offset():
    conv_offset = nn.Conv2d(
        inC,
        deformable_groups * 2 * kH * kW,
        kernel_size=(kH, kW),
        stride=(1, 1),
        padding=(1, 1),
        bias=True,
    ).cuda()

    conv_mask = nn.Conv2d(
        inC,
        deformable_groups * 1 * kH * kW,
        kernel_size=(kH, kW),
        stride=(1, 1),
        padding=(1, 1),
        bias=True,
    ).cuda()

    dcn_v2 = DCNv2(inC, outC, (kH, kW), stride=1, padding=1, dilation=1, deformable_groups=deformable_groups).cuda()
    scripted_dcnv2 = torch.jit.script(dcn_v2)

    conv_offset.weight.data.zero_()
    conv_offset.bias.data.zero_()
    conv_mask.weight.data.zero_()
    conv_mask.bias.data.zero_()
    conv_identify(dcn_v2.weight, dcn_v2.bias)

    input = torch.randn(N, inC, inH, inW).cuda()
    offset = conv_offset(input)
    mask = conv_mask(input)
    mask = torch.sigmoid(mask)
    output = dcn_v2(input, offset, mask)
    result = scripted_dcnv2(input, offset, mask)
    print(scripted_dcnv2.code)
    torch.testing.assert_allclose(output, result)

    output *= 2
    d = (input - output).abs().max()
    if d < 1e-10:
        print("Zero offset passed")
    else:
        print("Zero offset failed")
        print(input)
        print(output)
//  [python]
def check_zero_offset():
    conv_offset = nn.Conv2d(
        inC,
        deformable_groups * 2 * kH * kW,
        kernel_size=(kH, kW),
        stride=(1, 1),
        padding=(1, 1),
        bias=True,
    ).cuda()

    conv_mask = nn.Conv2d(
        inC,
        deformable_groups * 1 * kH * kW,
        kernel_size=(kH, kW),
        stride=(1, 1),
        padding=(1, 1),
        bias=True,
    ).cuda()

    dcn_v2 = DCNv2(inC, outC, (kH, kW), stride=1, padding=1, dilation=1, deformable_groups=deformable_groups).cuda()
    scripted_dcnv2 = torch.jit.script(dcn_v2)

    conv_offset.weight.data.zero_()
    conv_offset.bias.data.zero_()
    conv_mask.weight.data.zero_()
    conv_mask.bias.data.zero_()
    conv_identify(dcn_v2.weight, dcn_v2.bias)

    input = torch.randn(N, inC, inH, inW).cuda()
    offset = conv_offset(input)
    mask = conv_mask(input)
    mask = torch.sigmoid(mask)
    output = dcn_v2(input, offset, mask)
    result = scripted_dcnv2(input, offset, mask)
    print(scripted_dcnv2.code)
    torch.testing.assert_close(output, result)

    output *= 2
    d = (input - output).abs().max()
    if d < 1e-10:
        print("Zero offset passed")
    else:
        print("Zero offset failed")
        print(input)
        print(output)
//  [python]
def test_multi_output_gp(self):
        # Setup composite mean
        mean1 = gpytorch.means.ConstantMean()
        mean2 = gpytorch.means.ConstantMean()
        mean = BatchMean([mean1, mean2])

        # Setup composite kernel
        cov1 = gpytorch.kernels.RBFKernel()
        cov2 = gpytorch.kernels.RBFKernel()
        kernel = BatchKernel([cov1, cov2])

        # Training data
        train_x = torch.randn(5, 2)
        train_y = torch.randn(5, 2).t()

        # Combined GP
        likelihood = gpytorch.likelihoods.GaussianLikelihood(batch_size=2)
        gp = MultiOutputGP(train_x, train_y, kernel, likelihood, mean=mean)

        # Individual GPs
        likelihood1 = gpytorch.likelihoods.GaussianLikelihood()
        gp1 = ExactGPModel(train_x, train_y[0], cov1, likelihood1, mean=mean1)
        gp2 = ExactGPModel(train_x, train_y[1], cov2, likelihood1, mean=mean2)

        # Evaluation mode
        gp.eval()
        gp1.eval()
        gp2.eval()

        # Evaluate
        test_x = torch.randn(10, 2)
        pred = gp(test_x)
        pred1 = gp1(test_x)
        pred2 = gp2(test_x)

        torch.testing.assert_allclose(pred.mean[0], pred1.mean)
        torch.testing.assert_allclose(pred.mean[1], pred2.mean)

        torch.testing.assert_allclose(pred.covariance_matrix[0],
                                      pred1.covariance_matrix)
        torch.testing.assert_allclose(pred.covariance_matrix[1],
                                      pred2.covariance_matrix)

        torch.testing.assert_allclose(pred.variance[0], pred1.variance)
        torch.testing.assert_allclose(pred.variance[1], pred2.variance)

        # Test optimization
        gp.train()
        optimizer = torch.optim.Adam([{'params': gp.parameters()}], lr=0.1)
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp)
        optimizer.zero_grad()
        loss = gp.loss(mll)
        loss.backward()

        optimizer.step()
//  [python]
def test_multi_output_gp(self):
        # Setup composite mean
        mean1 = gpytorch.means.ConstantMean()
        mean2 = gpytorch.means.ConstantMean()
        mean = BatchMean([mean1, mean2])

        # Setup composite kernel
        cov1 = gpytorch.kernels.RBFKernel()
        cov2 = gpytorch.kernels.RBFKernel()
        kernel = BatchKernel([cov1, cov2])

        # Training data
        train_x = torch.randn(5, 2)
        train_y = torch.randn(5, 2).t()

        # Combined GP
        likelihood = gpytorch.likelihoods.GaussianLikelihood(batch_size=2)
        gp = MultiOutputGP(train_x, train_y, kernel, likelihood, mean=mean)

        # Individual GPs
        likelihood1 = gpytorch.likelihoods.GaussianLikelihood()
        gp1 = ExactGPModel(train_x, train_y[0], cov1, likelihood1, mean=mean1)
        gp2 = ExactGPModel(train_x, train_y[1], cov2, likelihood1, mean=mean2)

        # Evaluation mode
        gp.eval()
        gp1.eval()
        gp2.eval()

        # Evaluate
        test_x = torch.randn(10, 2)
        pred = gp(test_x)
        pred1 = gp1(test_x)
        pred2 = gp2(test_x)

        torch.testing.assert_close(pred.mean[0], pred1.mean)
        torch.testing.assert_close(pred.mean[1], pred2.mean)

        torch.testing.assert_close(pred.covariance_matrix[0],
                                      pred1.covariance_matrix)
        torch.testing.assert_close(pred.covariance_matrix[1],
                                      pred2.covariance_matrix)

        torch.testing.assert_close(pred.variance[0], pred1.variance)
        torch.testing.assert_close(pred.variance[1], pred2.variance)

        # Test optimization
        gp.train()
        optimizer = torch.optim.Adam([{'params': gp.parameters()}], lr=0.1)
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp)
        optimizer.zero_grad()
        loss = gp.loss(mll)
        loss.backward()

        optimizer.step()
//  [python]
def forward(
        self, ctx: Context, batch: Batch, grids: List[torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        for grid in grids:
            grid.mean().backward()

        ctx.add_scalar("test", 1234)

        for cam_feat in ctx.cam_feats.values():
            cam_feat.mean().backward()

        # check that start position is at zero
        car_to_world = batch.car_to_world(ctx.start_frame)
        zero = (
            car_to_world.new_tensor([0.0, 0.0, 0.0, 1.0])
            .expand(car_to_world.size(0), -1)
            .unsqueeze(-1)
        )
        long_cam_T, mask, lengths = batch.long_cam_T
        assert isinstance(long_cam_T, torch.Tensor)
        long_car_to_world = long_cam_T[:, ctx.start_frame].pinverse()
        torch.testing.assert_allclose(car_to_world.matmul(zero), zero)
        torch.testing.assert_allclose(long_car_to_world.matmul(zero), zero)

        return {
            "foo": torch.tensor(5.0),
        }
//  [python]
def forward(
        self, ctx: Context, batch: Batch, grids: List[torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        for grid in grids:
            grid.mean().backward()

        ctx.add_scalar("test", 1234)

        for cam_feat in ctx.cam_feats.values():
            cam_feat.mean().backward()

        # check that start position is at zero
        car_to_world = batch.car_to_world(ctx.start_frame)
        zero = (
            car_to_world.new_tensor([0.0, 0.0, 0.0, 1.0])
            .expand(car_to_world.size(0), -1)
            .unsqueeze(-1)
        )
        long_cam_T, mask, lengths = batch.long_cam_T
        assert isinstance(long_cam_T, torch.Tensor)
        long_car_to_world = long_cam_T[:, ctx.start_frame].pinverse()
        torch.testing.assert_close(car_to_world.matmul(zero), zero)
        torch.testing.assert_close(long_car_to_world.matmul(zero), zero)

        return {
            "foo": torch.tensor(5.0),
        }
//  [python]
def test_superpos(self):
        stack = MultiPopStack.empty(1, 3, None)
        stack.update(REDUCE0, NEW_VEC)
        policy = torch.tensor([[1/2, 1/2, 0, 0, 0, 0]])
        stack.update(policy, NEW_VEC)
        expected = [[[1., 1., 0.], [1/2, 1/2, 0]]]
        torch.testing.assert_allclose(stack.tapes.tolist(), expected)
//  [python]
def test_superpos(self):
        stack = MultiPopStack.empty(1, 3, None)
        stack.update(REDUCE0, NEW_VEC)
        policy = torch.tensor([[1/2, 1/2, 0, 0, 0, 0]])
        stack.update(policy, NEW_VEC)
        expected = [[[1., 1., 0.], [1/2, 1/2, 0]]]
        torch.testing.assert_close(stack.tapes.tolist(), expected)