jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_init(memory):
    torch.testing.assert_close(memory.next_done, torch.tensor([0, 0, 0, 0]))
    assert memory.next_value is None
    assert memory.device == torch.device("cpu")
    assert memory.global_step == 0
    assert memory.obs_preprocessor is not None
    assert memory.experiences == []
    assert memory.episode_lengths == []
    assert memory.episode_returns == []
    assert memory.vars_to_log == {}
//  [python]
def _test_rowwise_prune_op(self, embedding_rows, embedding_dims, indices_type, weights_dtype):
        embedding_weights = None
        if weights_dtype in [torch.int8, torch.int16, torch.int32, torch.int64]:
            embedding_weights = torch.randint(0, 100, (embedding_rows, embedding_dims), dtype=weights_dtype)
        else:
            embedding_weights = torch.rand((embedding_rows, embedding_dims), dtype=weights_dtype)
        mask = self._generate_rowwise_mask(embedding_rows)

        def get_pt_result(embedding_weights, mask, indices_type):
            return torch.rowwise_prune(embedding_weights, mask, indices_type)

        # Reference implementation.
        def get_reference_result(embedding_weights, mask, indices_type):
            num_embeddings = mask.size()[0]
            compressed_idx_out = torch.zeros(num_embeddings, dtype=indices_type)
            pruned_weights_out = embedding_weights[mask[:]]
            idx = 0
            for i in range(mask.size()[0]):
                if mask[i]:
                    compressed_idx_out[i] = idx
                    idx = idx + 1
                else:
                    compressed_idx_out[i] = -1
            return (pruned_weights_out, compressed_idx_out)

        pt_pruned_weights, pt_compressed_indices_map = get_pt_result(
            embedding_weights, mask, indices_type)
        ref_pruned_weights, ref_compressed_indices_map = get_reference_result(
            embedding_weights, mask, indices_type)

        torch.testing.assert_allclose(pt_pruned_weights, ref_pruned_weights)
        self.assertEqual(pt_compressed_indices_map, ref_compressed_indices_map)
        self.assertEqual(pt_compressed_indices_map.dtype, indices_type)
//  [python]
def _test_rowwise_prune_op(self, embedding_rows, embedding_dims, indices_type, weights_dtype):
        embedding_weights = None
        if weights_dtype in [torch.int8, torch.int16, torch.int32, torch.int64]:
            embedding_weights = torch.randint(0, 100, (embedding_rows, embedding_dims), dtype=weights_dtype)
        else:
            embedding_weights = torch.rand((embedding_rows, embedding_dims), dtype=weights_dtype)
        mask = self._generate_rowwise_mask(embedding_rows)

        def get_pt_result(embedding_weights, mask, indices_type):
            return torch.rowwise_prune(embedding_weights, mask, indices_type)

        # Reference implementation.
        def get_reference_result(embedding_weights, mask, indices_type):
            num_embeddings = mask.size()[0]
            compressed_idx_out = torch.zeros(num_embeddings, dtype=indices_type)
            pruned_weights_out = embedding_weights[mask[:]]
            idx = 0
            for i in range(mask.size()[0]):
                if mask[i]:
                    compressed_idx_out[i] = idx
                    idx = idx + 1
                else:
                    compressed_idx_out[i] = -1
            return (pruned_weights_out, compressed_idx_out)

        pt_pruned_weights, pt_compressed_indices_map = get_pt_result(
            embedding_weights, mask, indices_type)
        ref_pruned_weights, ref_compressed_indices_map = get_reference_result(
            embedding_weights, mask, indices_type)

        torch.testing.assert_close(pt_pruned_weights, ref_pruned_weights)
        self.assertEqual(pt_compressed_indices_map, ref_compressed_indices_map)
        self.assertEqual(pt_compressed_indices_map.dtype, indices_type)
//  [python]
def test_minibatch_class():
    obs = torch.tensor([1.0, 2.0, 3.0])
    actions = torch.tensor([1.0, 2.0, 3.0])
    logprobs = torch.tensor([1.0, 2.0, 3.0])
    advantages = torch.tensor([1.0, 2.0, 3.0])
    values = torch.tensor([1.0, 2.0, 3.0])
    returns = torch.tensor([1.0, 2.0, 3.0])

    minibatch = Minibatch(
        obs=obs,
        actions=actions,
        logprobs=logprobs,
        advantages=advantages,
        values=values,
        returns=returns,
    )

    torch.testing.assert_allclose(minibatch.obs, obs)
    torch.testing.assert_allclose(minibatch.actions, actions)
    torch.testing.assert_allclose(minibatch.logprobs, logprobs)
    torch.testing.assert_allclose(minibatch.advantages, advantages)
    torch.testing.assert_allclose(minibatch.values, values)
    torch.testing.assert_allclose(minibatch.returns, returns)
//  [python]
def test_minibatch_class():
    obs = torch.tensor([1.0, 2.0, 3.0])
    actions = torch.tensor([1.0, 2.0, 3.0])
    logprobs = torch.tensor([1.0, 2.0, 3.0])
    advantages = torch.tensor([1.0, 2.0, 3.0])
    values = torch.tensor([1.0, 2.0, 3.0])
    returns = torch.tensor([1.0, 2.0, 3.0])

    minibatch = Minibatch(
        obs=obs,
        actions=actions,
        logprobs=logprobs,
        advantages=advantages,
        values=values,
        returns=returns,
    )

    torch.testing.assert_close(minibatch.obs, obs)
    torch.testing.assert_close(minibatch.actions, actions)
    torch.testing.assert_close(minibatch.logprobs, logprobs)
    torch.testing.assert_close(minibatch.advantages, advantages)
    torch.testing.assert_close(minibatch.values, values)
    torch.testing.assert_close(minibatch.returns, returns)
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [DataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
        torch.testing.assert_allclose(data_samples[0].pred_score,
                                      predictions[0].pred_score)
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [DataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_score.shape, (10, ))
        torch.testing.assert_close(data_samples[0].pred_score,
                                      predictions[0].pred_score)
//  [python]
def test_deit_tiny():
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    model_config = [
        ("BACKEND", "tensorrt"),
        ("SCHEDULE.FUSE_BN", True),
        ("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"),
        ("A.OBSERVER.LAYOUT", "NCHW"),
        ("A.QADD.ENABLE_QUANT", True),
    ]
    model_config = [j for i in model_config for j in i]

    model = ConvAdd()
    config = build_config(model_config)
    qmodel = QuantModel(model, config)

    data = torch.randn(1, 3, 4, 4)
    model.eval()
    qmodel.eval()
    out1 = model(data)
    out2 = qmodel(data)
    for node in qmodel.model.graph.nodes:
        module = getattr(qmodel.model, node.target, None)
        if module is not None:
            if isinstance(module, QAdd):
                counter = 0
                for prev_node in node.all_input_nodes:
                    input_module = getattr(qmodel.model, prev_node.target, None)
                    if isinstance(input_module, QIdentity):
                        counter += 1
                assert counter > 0, "Qadd quantizers not build successfully!"
    torch.testing.assert_allclose(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test_deit_tiny():
    # the format of list([k1, v1, k2, v2, ...]), ensure config[::2] is key and config[1::2] is value
    model_config = [
        ("BACKEND", "tensorrt"),
        ("SCHEDULE.FUSE_BN", True),
        ("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"),
        ("A.OBSERVER.LAYOUT", "NCHW"),
        ("A.QADD.ENABLE_QUANT", True),
    ]
    model_config = [j for i in model_config for j in i]

    model = ConvAdd()
    config = build_config(model_config)
    qmodel = QuantModel(model, config)

    data = torch.randn(1, 3, 4, 4)
    model.eval()
    qmodel.eval()
    out1 = model(data)
    out2 = qmodel(data)
    for node in qmodel.model.graph.nodes:
        module = getattr(qmodel.model, node.target, None)
        if module is not None:
            if isinstance(module, QAdd):
                counter = 0
                for prev_node in node.all_input_nodes:
                    input_module = getattr(qmodel.model, prev_node.target, None)
                    if isinstance(input_module, QIdentity):
                        counter += 1
                assert counter > 0, "Qadd quantizers not build successfully!"
    torch.testing.assert_close(out1, out2, atol=1e-4, rtol=1e-4)
//  [python]
def test_compute_advantages(memory):
    info = {"final_info": [{"episode": {"l": 1, "r": 1.0}}]}
    obs = torch.tensor([1.0, 2.0, 3.0])
    done = torch.tensor([1.0, 2.0, 3.0])
    action = torch.tensor([1.0, 2.0, 3.0])
    logprob = torch.tensor([1.0, 2.0, 3.0])
    value = torch.tensor([1.0, 2.0, 3.0])
    reward = torch.tensor([1.0, 2.0, 3.0])

    memory.add(info, obs, done, action, logprob, value, reward)

    advantages = memory.compute_advantages(
        next_value=torch.tensor([1.0, 2.0, 3.0]),
        next_done=torch.tensor([0, 1, 0]),
        rewards=torch.tensor([1.0, 2.0, 3.0]).repeat(2, 1),
        values=torch.tensor([1.0, 2.0, 3.0]).repeat(2, 1),
        dones=torch.tensor([1, 0, 0]).repeat(2, 1),
        device=torch.device("cpu"),
        gamma=0.99,
        gae_lambda=0.95,
    )

    torch.testing.assert_allclose(
        advantages, torch.tensor([[0.0, 1.98, 5.7633], [0.99, 0.00, 2.97]])
    )
//  [python]
def test_compute_advantages(memory):
    info = {"final_info": [{"episode": {"l": 1, "r": 1.0}}]}
    obs = torch.tensor([1.0, 2.0, 3.0])
    done = torch.tensor([1.0, 2.0, 3.0])
    action = torch.tensor([1.0, 2.0, 3.0])
    logprob = torch.tensor([1.0, 2.0, 3.0])
    value = torch.tensor([1.0, 2.0, 3.0])
    reward = torch.tensor([1.0, 2.0, 3.0])

    memory.add(info, obs, done, action, logprob, value, reward)

    advantages = memory.compute_advantages(
        next_value=torch.tensor([1.0, 2.0, 3.0]),
        next_done=torch.tensor([0, 1, 0]),
        rewards=torch.tensor([1.0, 2.0, 3.0]).repeat(2, 1),
        values=torch.tensor([1.0, 2.0, 3.0]).repeat(2, 1),
        dones=torch.tensor([1, 0, 0]).repeat(2, 1),
        device=torch.device("cpu"),
        gamma=0.99,
        gae_lambda=0.95,
    )

    torch.testing.assert_close(
        advantages, torch.tensor([[0.0, 1.98, 5.7633], [0.99, 0.00, 2.97]])
    )
//  [python]
def test_test_step(self):
        cfg = ConfigDict(deepcopy(self.DEFAULT_ARGS))
        cfg.module.data_preprocessor = dict(
            mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5])
        model: AverageClsScoreTTA = MODELS.build(cfg)

        img1 = torch.randint(0, 256, (1, 3, 224, 224))
        img2 = torch.randint(0, 256, (1, 3, 224, 224))
        data1 = {
            'inputs': img1,
            'data_samples': [DataSample().set_gt_label(1)]
        }
        data2 = {
            'inputs': img2,
            'data_samples': [DataSample().set_gt_label(1)]
        }
        data_tta = {
            'inputs': [img1, img2],
            'data_samples': [[DataSample().set_gt_label(1)],
                             [DataSample().set_gt_label(1)]]
        }

        score1 = model.module.test_step(data1)[0].pred_score
        score2 = model.module.test_step(data2)[0].pred_score
        score_tta = model.test_step(data_tta)[0].pred_score

        torch.testing.assert_allclose(score_tta, (score1 + score2) / 2)
//  [python]
def test_test_step(self):
        cfg = ConfigDict(deepcopy(self.DEFAULT_ARGS))
        cfg.module.data_preprocessor = dict(
            mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5])
        model: AverageClsScoreTTA = MODELS.build(cfg)

        img1 = torch.randint(0, 256, (1, 3, 224, 224))
        img2 = torch.randint(0, 256, (1, 3, 224, 224))
        data1 = {
            'inputs': img1,
            'data_samples': [DataSample().set_gt_label(1)]
        }
        data2 = {
            'inputs': img2,
            'data_samples': [DataSample().set_gt_label(1)]
        }
        data_tta = {
            'inputs': [img1, img2],
            'data_samples': [[DataSample().set_gt_label(1)],
                             [DataSample().set_gt_label(1)]]
        }

        score1 = model.module.test_step(data1)[0].pred_score
        score2 = model.module.test_step(data2)[0].pred_score
        score_tta = model.test_step(data_tta)[0].pred_score

        torch.testing.assert_close(score_tta, (score1 + score2) / 2)
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [DataSample().set_gt_label(1)]
        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

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

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

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

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

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

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

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

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

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [DataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_score.shape, (1000, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_score.shape, (1000, ))
        self.assertEqual(data_samples[0].pred_score.shape, (1000, ))
        torch.testing.assert_allclose(data_samples[0].pred_score,
                                      predictions[0].pred_score)
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [DataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_score.shape, (1000, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_score.shape, (1000, ))
        self.assertEqual(data_samples[0].pred_score.shape, (1000, ))
        torch.testing.assert_close(data_samples[0].pred_score,
                                      predictions[0].pred_score)
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [DataSample().set_gt_label(1)]
        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

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

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

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

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

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

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

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

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

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_bdd100ksemseg(self) -> None:
        device = torch.device("cpu")
        mmlab = BDD100KSemSeg(device, mmlab=True)
        ts = BDD100KSemSeg(device, mmlab=False)
        inp = torch.rand(1, 3, 48, 64, device=device)
        out_mmlab = mmlab(inp)
        out_ts = ts(inp)

        torch.testing.assert_allclose(out_mmlab, out_ts)
//  [python]
def test_bdd100ksemseg(self) -> None:
        device = torch.device("cpu")
        mmlab = BDD100KSemSeg(device, mmlab=True)
        ts = BDD100KSemSeg(device, mmlab=False)
        inp = torch.rand(1, 3, 48, 64, device=device)
        out_mmlab = mmlab(inp)
        out_ts = ts(inp)

        torch.testing.assert_close(out_mmlab, out_ts)
//  [python]
def test_batch_bucket_one_hot_op(self):
        data = np.array([[2, 3], [4, 1], [2, 5]]).astype(np.float32)
        lengths = np.array([2, 3]).astype(np.int32)
        boundaries = np.array([0.1, 2.5, 1, 3.1, 4.5]).astype(np.float32)

        def _batch_bucket_one_hot_ref(data, lengths, boundaries):
            ref_op = core.CreateOperator(
                "BatchBucketOneHot", ["data", "lengths", "boundaries"], ["Y"]
            )
            workspace.FeedBlob("data", data)
            workspace.FeedBlob("lengths", lengths)
            workspace.FeedBlob("boundaries", boundaries)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _batch_bucket_one_hot_ref(data, lengths, boundaries)
        actual_output = torch.ops._caffe2.BatchBucketOneHot(
            torch.tensor(data), torch.Tensor(lengths).int(), torch.Tensor(boundaries)
        )
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_batch_bucket_one_hot_op(self):
        data = np.array([[2, 3], [4, 1], [2, 5]]).astype(np.float32)
        lengths = np.array([2, 3]).astype(np.int32)
        boundaries = np.array([0.1, 2.5, 1, 3.1, 4.5]).astype(np.float32)

        def _batch_bucket_one_hot_ref(data, lengths, boundaries):
            ref_op = core.CreateOperator(
                "BatchBucketOneHot", ["data", "lengths", "boundaries"], ["Y"]
            )
            workspace.FeedBlob("data", data)
            workspace.FeedBlob("lengths", lengths)
            workspace.FeedBlob("boundaries", boundaries)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _batch_bucket_one_hot_ref(data, lengths, boundaries)
        actual_output = torch.ops._caffe2.BatchBucketOneHot(
            torch.tensor(data), torch.Tensor(lengths).int(), torch.Tensor(boundaries)
        )
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_percentile(self):
        original_values = np.array([[3.0, 5.0, 3], [5.0, 1.0, 6.0]]).astype(np.float32)
        value_to_pct = np.array([[3, 0.2], [5, 0.5], [1, 0.3], [3, 0.6]]).astype(
            np.float32
        )
        lengths = np.array([2, 1, 1]).astype(np.int32)

        def _percentile_ref(original_values, value_to_pct, lengths):
            ref_op = core.CreateOperator(
                "Percentile", ["original_values", "value_to_pct", "lengths"], ["Y"]
            )
            workspace.FeedBlob("original_values", original_values)
            workspace.FeedBlob("value_to_pct", value_to_pct)
            workspace.FeedBlob("lengths", lengths)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _percentile_ref(original_values, value_to_pct, lengths)
        actual_output = torch.ops._caffe2.Percentile(
            torch.tensor(original_values),
            torch.Tensor(value_to_pct),
            torch.Tensor(lengths).int(),
        )
        torch.testing.assert_allclose(expected_output, actual_output.cpu())
//  [python]
def test_percentile(self):
        original_values = np.array([[3.0, 5.0, 3], [5.0, 1.0, 6.0]]).astype(np.float32)
        value_to_pct = np.array([[3, 0.2], [5, 0.5], [1, 0.3], [3, 0.6]]).astype(
            np.float32
        )
        lengths = np.array([2, 1, 1]).astype(np.int32)

        def _percentile_ref(original_values, value_to_pct, lengths):
            ref_op = core.CreateOperator(
                "Percentile", ["original_values", "value_to_pct", "lengths"], ["Y"]
            )
            workspace.FeedBlob("original_values", original_values)
            workspace.FeedBlob("value_to_pct", value_to_pct)
            workspace.FeedBlob("lengths", lengths)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("Y")

        expected_output = _percentile_ref(original_values, value_to_pct, lengths)
        actual_output = torch.ops._caffe2.Percentile(
            torch.tensor(original_values),
            torch.Tensor(value_to_pct),
            torch.Tensor(lengths).int(),
        )
        torch.testing.assert_close(expected_output, actual_output.cpu())
//  [python]
def test_from_nir_to_sequential():
    batch_size = 4

    orig_model = nn.Sequential(
        torch.nn.Linear(10, 2),
        sl.ExpLeakSqueeze(tau_mem=10.0, batch_size=batch_size),
        sl.LIFSqueeze(tau_mem=10.0, batch_size=batch_size),
        torch.nn.Linear(2, 1),
    )
    nir_graph = to_nir(orig_model, torch.randn(batch_size, 10))

    converted_model = from_nir(nir_graph, batch_size=batch_size)
    converted_modules = list(converted_model.children())
    assert len(orig_model) + 2 == len(
        converted_modules
    )  # Addition of input and output modules
    torch.testing.assert_allclose(orig_model[0].weight, converted_modules[1].weight)
    torch.testing.assert_allclose(orig_model[0].bias, converted_modules[1].bias)
    assert type(orig_model[1]) == type(converted_modules[2])
    assert type(orig_model[2]) == type(converted_modules[3])
    torch.testing.assert_allclose(orig_model[3].weight, converted_modules[4].weight)
    torch.testing.assert_allclose(orig_model[3].bias, converted_modules[4].bias)
//  [python]
def test_from_nir_to_sequential():
    batch_size = 4

    orig_model = nn.Sequential(
        torch.nn.Linear(10, 2),
        sl.ExpLeakSqueeze(tau_mem=10.0, batch_size=batch_size),
        sl.LIFSqueeze(tau_mem=10.0, batch_size=batch_size),
        torch.nn.Linear(2, 1),
    )
    nir_graph = to_nir(orig_model, torch.randn(batch_size, 10))

    converted_model = from_nir(nir_graph, batch_size=batch_size)
    converted_modules = list(converted_model.children())
    assert len(orig_model) + 2 == len(
        converted_modules
    )  # Addition of input and output modules
    torch.testing.assert_close(orig_model[0].weight, converted_modules[1].weight)
    torch.testing.assert_close(orig_model[0].bias, converted_modules[1].bias)
    assert type(orig_model[1]) == type(converted_modules[2])
    assert type(orig_model[2]) == type(converted_modules[3])
    torch.testing.assert_close(orig_model[3].weight, converted_modules[4].weight)
    torch.testing.assert_close(orig_model[3].bias, converted_modules[4].bias)
//  [python]
def _test_roi_align_rotated(self, N, C, H, W, device):
        def rand_rotated_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * 360 - 180,
                ]
            ).astype(np.float32)

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

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlignRotated",
                ["feature", "rois"],
                ["roi_feature"],
                spatial_scale=1.0,
                pooled_h=3,
                pooled_w=3,
                sampling_ratio=0,
            )
            workspace.FeedBlob("feature", _feature)
            workspace.FeedBlob("rois", _rois)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("roi_feature")

        roi_feature_ref = roi_align_ref(feature, rois)
        roi_feature = torch.ops._caffe2.RoIAlignRotated(
            torch.Tensor(feature).to(device),
            torch.Tensor(rois).to(device),
            order="NCHW",
            spatial_scale=1.0,
            pooled_h=3,
            pooled_w=3,
            sampling_ratio=0,
            aligned=False,
        )
        torch.testing.assert_allclose(roi_feature_ref, roi_feature.cpu())
//  [python]
def _test_roi_align_rotated(self, N, C, H, W, device):
        def rand_rotated_roi():
            return np.array(
                [
                    float(int(N * np.random.rand())),
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * W,
                    np.random.rand() * H,
                    np.random.rand() * 360 - 180,
                ]
            ).astype(np.float32)

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

        def roi_align_ref(_feature, _rois):
            ref_op = core.CreateOperator(
                "RoIAlignRotated",
                ["feature", "rois"],
                ["roi_feature"],
                spatial_scale=1.0,
                pooled_h=3,
                pooled_w=3,
                sampling_ratio=0,
            )
            workspace.FeedBlob("feature", _feature)
            workspace.FeedBlob("rois", _rois)
            workspace.RunOperatorOnce(ref_op)
            return workspace.FetchBlob("roi_feature")

        roi_feature_ref = roi_align_ref(feature, rois)
        roi_feature = torch.ops._caffe2.RoIAlignRotated(
            torch.Tensor(feature).to(device),
            torch.Tensor(rois).to(device),
            order="NCHW",
            spatial_scale=1.0,
            pooled_h=3,
            pooled_w=3,
            sampling_ratio=0,
            aligned=False,
        )
        torch.testing.assert_close(roi_feature_ref, roi_feature.cpu())
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_rms_layers(self.normalized_shape, elementwise_affine, dtype)
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x.cpu())
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedRMSNorm.bf16_fwd_thresholds
        torch.testing.assert_allclose(actual, expected.detach().clone().cuda(), **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.detach().clone().cuda()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedRMSNorm.bf16_bwd_thresholds
        torch.testing.assert_allclose(native_x.grad.cuda(), fused_x.grad, **tols)
//  [python]
def _run_test(self, dtype, elementwise_affine):
        native, fused = _prep_rms_layers(self.normalized_shape, elementwise_affine, dtype)
        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)

        expected = native(native_x.cpu())
        with torch.cuda.amp.autocast(dtype=dtype):
            actual = fused(fused_x)
        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedRMSNorm.bf16_fwd_thresholds
        torch.testing.assert_close(actual, expected.detach().clone().cuda(), **tols)

        g_native = torch.rand_like(expected)
        with torch.no_grad():
            g_fused = g_native.detach().clone().cuda()
        expected.backward(g_native)
        actual.backward(g_fused)

        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedRMSNorm.bf16_bwd_thresholds
        torch.testing.assert_close(native_x.grad.cuda(), fused_x.grad, **tols)
//  [python]
def test_scalar_invariance():
    torch.manual_seed(1234)
    rotate = torch.tensor(
        [
            [0.9886788, -0.1102370, 0.1017945],
            [0.1363630, 0.9431761, -0.3030248],
            [-0.0626055, 0.3134752, 0.9475304],
        ]
    )

    model = create_model(load_example_args("equivariant-transformer", prior_model=None))
    z = torch.ones(100, dtype=torch.long)
    pos = torch.randn(100, 3)
    batch = torch.arange(50, dtype=torch.long).repeat_interleave(2)

    y = model(z, pos, batch)[0]
    y_rot = model(z, pos @ rotate, batch)[0]
    torch.testing.assert_allclose(y, y_rot)
//  [python]
def test_scalar_invariance():
    torch.manual_seed(1234)
    rotate = torch.tensor(
        [
            [0.9886788, -0.1102370, 0.1017945],
            [0.1363630, 0.9431761, -0.3030248],
            [-0.0626055, 0.3134752, 0.9475304],
        ]
    )

    model = create_model(load_example_args("equivariant-transformer", prior_model=None))
    z = torch.ones(100, dtype=torch.long)
    pos = torch.randn(100, 3)
    batch = torch.arange(50, dtype=torch.long).repeat_interleave(2)

    y = model(z, pos, batch)[0]
    y_rot = model(z, pos @ rotate, batch)[0]
    torch.testing.assert_close(y, y_rot)
//  [python]
def test_vector_equivariance():
    torch.manual_seed(1234)
    rotate = torch.tensor(
        [
            [0.9886788, -0.1102370, 0.1017945],
            [0.1363630, 0.9431761, -0.3030248],
            [-0.0626055, 0.3134752, 0.9475304],
        ]
    )

    model = create_model(
        load_example_args(
            "equivariant-transformer",
            prior_model=None,
            output_model="VectorOutput",
        )
    )
    z = torch.ones(100, dtype=torch.long)
    pos = torch.randn(100, 3)
    batch = torch.arange(50, dtype=torch.long).repeat_interleave(2)

    y = model(z, pos, batch)[0]
    y_rot = model(z, pos @ rotate, batch)[0]
    torch.testing.assert_allclose(y @ rotate, y_rot)
//  [python]
def test_vector_equivariance():
    torch.manual_seed(1234)
    rotate = torch.tensor(
        [
            [0.9886788, -0.1102370, 0.1017945],
            [0.1363630, 0.9431761, -0.3030248],
            [-0.0626055, 0.3134752, 0.9475304],
        ]
    )

    model = create_model(
        load_example_args(
            "equivariant-transformer",
            prior_model=None,
            output_model="VectorOutput",
        )
    )
    z = torch.ones(100, dtype=torch.long)
    pos = torch.randn(100, 3)
    batch = torch.arange(50, dtype=torch.long).repeat_interleave(2)

    y = model(z, pos, batch)[0]
    y_rot = model(z, pos @ rotate, batch)[0]
    torch.testing.assert_close(y @ rotate, y_rot)
//  [python]
def test_v_trace():
    ratios = torch.tensor([1., 1., 1.])
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

    reset = torch.tensor([False, False, False])
    actual = v_trace(ratios, value, reward, reset, gamma)
    torch.testing.assert_allclose(actual, torch.tensor([11., 9., 6.]))

    reset = torch.tensor([False, True, False])
    actual = v_trace(ratios, value, reward, reset, gamma)
    torch.testing.assert_allclose(actual, torch.tensor([2., 9., 6.]))
//  [python]
def test_v_trace():
    ratios = torch.tensor([1., 1., 1.])
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

    reset = torch.tensor([False, False, False])
    actual = v_trace(ratios, value, reward, reset, gamma)
    torch.testing.assert_close(actual, torch.tensor([11., 9., 6.]))

    reset = torch.tensor([False, True, False])
    actual = v_trace(ratios, value, reward, reset, gamma)
    torch.testing.assert_close(actual, torch.tensor([2., 9., 6.]))
//  [python]
def test_v_trace_ref():
    ratios = torch.tensor([1., 1., 1.])
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

    reset = torch.tensor([False, False, False])
    actual = v_trace_ref(ratios, value, reward, reset, gamma)
    torch.testing.assert_allclose(actual, torch.tensor([11., 9., 6.]))

    reset = torch.tensor([False, True, False])
    actual = v_trace_ref(ratios, value, reward, reset, gamma)
    torch.testing.assert_allclose(actual, torch.tensor([2., 9., 6.]))
//  [python]
def test_v_trace_ref():
    ratios = torch.tensor([1., 1., 1.])
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

    reset = torch.tensor([False, False, False])
    actual = v_trace_ref(ratios, value, reward, reset, gamma)
    torch.testing.assert_close(actual, torch.tensor([11., 9., 6.]))

    reset = torch.tensor([False, True, False])
    actual = v_trace_ref(ratios, value, reward, reset, gamma)
    torch.testing.assert_close(actual, torch.tensor([2., 9., 6.]))
//  [python]
def _check_same_output(self, batch_size, contiguous):
        torch.cuda.manual_seed(42)
        if contiguous:
            input_shape = [batch_size] + self.normalized_shape
            input_ = torch.randn(input_shape, device="cpu").requires_grad_(True)
            input_cuda_ = input_.to(device="cuda", dtype=self.dtype).detach().requires_grad_(True)
            self.assertTrue(input_.is_contiguous())
            self.assertTrue(input_cuda_.is_contiguous())
        else:
            input_shape = [batch_size] + self.normalized_shape
            input_shape = [batch_size * 3] + [self.normalized_shape[0] * 5, self.normalized_shape[1] * 3]
            input_src_ = torch.randn(input_shape, device="cpu")
            input_ = input_src_[::3, ::5, ::3].detach().requires_grad_(True)
            input_cuda_ = input_src_.to(device="cuda", dtype=self.dtype)[::3, ::5, ::3].detach().requires_grad_(True)
            # make sure that tensors are NOT contiguous.
            self.assertFalse(input_.is_contiguous())
            self.assertFalse(input_cuda_.is_contiguous())
        out_cpu_ = self.module_cpu_(input_)
        gO = torch.rand_like(out_cpu_)
        out_cpu_.backward(gO)
        out_cuda_ = self.module_cuda_(input_cuda_)
        # TODO (mkozuki): `torch.testing.assert_allclose` is deprecated.
        # Use `torch.testing.assert_close`.
        # See https://github.com/pytorch/pytorch/issues/61844
        torch.testing.assert_allclose(
            out_cpu_.to(device="cuda", dtype=self.dtype), out_cuda_.clone().detach(), **self.fwd_thresholds)
        gO = gO.to(device="cuda", dtype=self.dtype)
        out_cuda_.backward(gO)
        self.assertFalse(out_cpu_.is_cuda)
        self.assertTrue(out_cuda_.is_cuda)
        torch.testing.assert_allclose(
            input_.grad.to(device="cuda", dtype=self.dtype), input_cuda_.grad, **self.bwd_thresholds)
        if self.elementwise_affine:
            torch.testing.assert_allclose(self.module_cpu_.weight.grad.to(device="cuda", dtype=self.dtype),
                                          self.module_cuda_.weight.grad, **self.bwd_thresholds)
//  [python]
def _check_same_output(self, batch_size, contiguous):
        torch.cuda.manual_seed(42)
        if contiguous:
            input_shape = [batch_size] + self.normalized_shape
            input_ = torch.randn(input_shape, device="cpu").requires_grad_(True)
            input_cuda_ = input_.to(device="cuda", dtype=self.dtype).detach().requires_grad_(True)
            self.assertTrue(input_.is_contiguous())
            self.assertTrue(input_cuda_.is_contiguous())
        else:
            input_shape = [batch_size] + self.normalized_shape
            input_shape = [batch_size * 3] + [self.normalized_shape[0] * 5, self.normalized_shape[1] * 3]
            input_src_ = torch.randn(input_shape, device="cpu")
            input_ = input_src_[::3, ::5, ::3].detach().requires_grad_(True)
            input_cuda_ = input_src_.to(device="cuda", dtype=self.dtype)[::3, ::5, ::3].detach().requires_grad_(True)
            # make sure that tensors are NOT contiguous.
            self.assertFalse(input_.is_contiguous())
            self.assertFalse(input_cuda_.is_contiguous())
        out_cpu_ = self.module_cpu_(input_)
        gO = torch.rand_like(out_cpu_)
        out_cpu_.backward(gO)
        out_cuda_ = self.module_cuda_(input_cuda_)
        torch.testing.assert_close(
            out_cpu_.to(device="cuda", dtype=self.dtype), out_cuda_.clone().detach(), **self.fwd_thresholds)
        gO = gO.to(device="cuda", dtype=self.dtype)
        out_cuda_.backward(gO)
        self.assertFalse(out_cpu_.is_cuda)
        self.assertTrue(out_cuda_.is_cuda)
        torch.testing.assert_close(
            input_.grad.to(device="cuda", dtype=self.dtype), input_cuda_.grad, **self.bwd_thresholds)
        if self.elementwise_affine:
            torch.testing.assert_close(self.module_cpu_.weight.grad.to(device="cuda", dtype=self.dtype),
                                          self.module_cuda_.weight.grad, **self.bwd_thresholds)
//  [python]
def test_v_trace_equivalent(R=100, T=10):
    for _ in range(R):
        ratios = torch.rand((T,))
        value = torch.rand((T,))
        reward = torch.rand((T,))
        reset = torch.rand((T,)) > .8
        gamma = torch.rand(())

        expected = v_trace_ref(ratios, value, reward, reset, gamma)
        actual = v_trace(ratios, value, reward, reset, gamma)

        torch.testing.assert_allclose(expected, actual)
//  [python]
def test_v_trace_equivalent(R=100, T=10):
    for _ in range(R):
        ratios = torch.rand((T,))
        value = torch.rand((T,))
        reward = torch.rand((T,))
        reset = torch.rand((T,)) > .8
        gamma = torch.rand(())

        expected = v_trace_ref(ratios, value, reward, reset, gamma)
        actual = v_trace(ratios, value, reward, reset, gamma)

        torch.testing.assert_close(expected, actual)
//  [python]
def test_forward(self):
        ga_covar = self.ga_kernel(self.X_in)
        ga_mean = delazify(self.ga_mean(self.X_in))
        ga_covar_diag = self.ga_kernel(self.X_in).diag()

        self.assertIsInstance(ga_covar, LazyEvaluatedKernelTensor)
        ga_covar = delazify(ga_covar)
        self.fitc.eval()

        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            fitc_output = self.fitc(self.X_in)
            fitc_mean = fitc_output.mean
            fitc_covar = fitc_output.covariance_matrix

        torch.testing.assert_allclose(fitc_mean, ga_mean)
        torch.testing.assert_allclose(ga_covar, fitc_covar)
        torch.testing.assert_allclose(ga_covar_diag, fitc_output.variance)

        with self.assertRaises(RuntimeError):
            delazify(self.ga_kernel(self.X_in, self.X_out))

        self.ga_kernel.eval()
        ga_covar_train_test = delazify(self.ga_kernel(self.X_in, self.X_out))

        self.assertTrue(ga_covar_train_test.shape == (len(self.X_in), len(self.X_out)))
//  [python]
def test_forward(self):
        ga_covar = self.ga_kernel(self.X_in)
        ga_mean = delazify(self.ga_mean(self.X_in))
        ga_covar_diag = self.ga_kernel(self.X_in).diag()

        self.assertIsInstance(ga_covar, LazyEvaluatedKernelTensor)
        ga_covar = delazify(ga_covar)
        self.fitc.eval()

        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            fitc_output = self.fitc(self.X_in)
            fitc_mean = fitc_output.mean
            fitc_covar = fitc_output.covariance_matrix

        torch.testing.assert_close(fitc_mean, ga_mean)
        torch.testing.assert_close(ga_covar, fitc_covar)
        torch.testing.assert_close(ga_covar_diag, fitc_output.variance)

        with self.assertRaises(RuntimeError):
            delazify(self.ga_kernel(self.X_in, self.X_out))

        self.ga_kernel.eval()
        ga_covar_train_test = delazify(self.ga_kernel(self.X_in, self.X_out))

        self.assertTrue(ga_covar_train_test.shape == (len(self.X_in), len(self.X_out)))
//  [python]
def test_cache(self):
        self.ga_kernel.train()
        delazify(self.ga_kernel(self.X_in))
        for cache in self.fitc_eval_cache.keys():
            self.assertFalse(hasattr(self.ga_kernel, cache))
        # Make sure that all the cached values are successfully stored
        for cache, value in self.fitc_train_cache.items():
            self.assertTrue(hasattr(self.ga_kernel, cache))
            if cache == "_train_cached_posterior_mean":
                torch.testing.assert_allclose(getattr(self.ga_kernel, cache), getattr(self.ga_kernel, value)(self.X_in))
            else:
                torch.testing.assert_allclose(getattr(self.ga_kernel, cache), getattr(self.ga_kernel, value))

        self.ga_kernel.eval()
        delazify(self.ga_kernel(self.X_in))
        delazify(self.ga_mean(self.X_in))
        for cache, value in self.fitc_eval_cache.items():
            self.assertTrue(hasattr(self.ga_kernel, cache))
            self.assertTrue(getattr(self.ga_kernel, cache) is getattr(self.ga_kernel, value))
        for cache in self.fitc_train_cache:
            self.assertFalse(hasattr(self.ga_kernel, cache))

        self.ga_kernel.train()
        for cache in self.fitc_eval_cache:
            self.assertFalse(hasattr(self.ga_kernel, cache))
        for cache in self.fitc_train_cache:
            self.assertFalse(hasattr(self.ga_kernel, cache))
//  [python]
def test_cache(self):
        self.ga_kernel.train()
        delazify(self.ga_kernel(self.X_in))
        for cache in self.fitc_eval_cache.keys():
            self.assertFalse(hasattr(self.ga_kernel, cache))
        # Make sure that all the cached values are successfully stored
        for cache, value in self.fitc_train_cache.items():
            self.assertTrue(hasattr(self.ga_kernel, cache))
            if cache == "_train_cached_posterior_mean":
                torch.testing.assert_close(getattr(self.ga_kernel, cache), getattr(self.ga_kernel, value)(self.X_in))
            else:
                torch.testing.assert_close(getattr(self.ga_kernel, cache), getattr(self.ga_kernel, value))

        self.ga_kernel.eval()
        delazify(self.ga_kernel(self.X_in))
        delazify(self.ga_mean(self.X_in))
        for cache, value in self.fitc_eval_cache.items():
            self.assertTrue(hasattr(self.ga_kernel, cache))
            self.assertTrue(getattr(self.ga_kernel, cache) is getattr(self.ga_kernel, value))
        for cache in self.fitc_train_cache:
            self.assertFalse(hasattr(self.ga_kernel, cache))

        self.ga_kernel.train()
        for cache in self.fitc_eval_cache:
            self.assertFalse(hasattr(self.ga_kernel, cache))
        for cache in self.fitc_train_cache:
            self.assertFalse(hasattr(self.ga_kernel, cache))
//  [python]
def test_reward_to_go():
    reward = torch.tensor([1., 2., 3.])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.

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

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

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

    reset = torch.tensor([False, True, False])
    actual = reward_to_go(reward, value, reset, gamma)
    torch.testing.assert_close(actual, torch.tensor([2., 9., 6.]))
//  [python]
def test_generalized_advantages():
    reward = torch.tensor([1., 2., 3])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.
    lambd = 1.

    reset = torch.tensor([False, False, False])
    adv = generalized_advantages(value, reward, value, reset, gamma=gamma, lambd=lambd)
    torch.testing.assert_allclose(adv, torch.tensor([7., 4., 0.]))

    reset = torch.tensor([False, True, False])
    adv = generalized_advantages(value, reward, value, reset, gamma=gamma, lambd=lambd)
    torch.testing.assert_allclose(adv, torch.tensor([-2., 4., 0.]))
//  [python]
def test_generalized_advantages():
    reward = torch.tensor([1., 2., 3])
    value = torch.tensor([4., 5., 6.])
    gamma = 1.
    lambd = 1.

    reset = torch.tensor([False, False, False])
    adv = generalized_advantages(value, reward, value, reset, gamma=gamma, lambd=lambd)
    torch.testing.assert_close(adv, torch.tensor([7., 4., 0.]))

    reset = torch.tensor([False, True, False])
    adv = generalized_advantages(value, reward, value, reset, gamma=gamma, lambd=lambd)
    torch.testing.assert_close(adv, torch.tensor([-2., 4., 0.]))
//  [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_allclose(outs[:, 1:, :], outs_with_mask[:, 1:, :])