jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_leaf_operator_reg(self):
        class Leaf(torch.nn.Module):
            def forward(self, x, y):
                return x + y + torch.nn.Parameter(torch.ones(5))

        leaf = Leaf()
        wrap_lib.define("wrapped_foo(Tensor x, Tensor y) -> Tensor")
        wrap_lib.impl("wrapped_foo", leaf, "CPU")

        class Bar(torch.nn.Module):
            def __init__(self):
                super(Bar, self).__init__()
                self.foo = torch.ops.wrap.wrapped_foo
                self.other = torch.nn.Parameter(torch.ones(5))

            def forward(self, x, y):
                x = self.foo(x, y)
                x = x + self.other
                return x

        mod = Bar().eval()
        inputs = [torch.ones(5), torch.ones(5)]
        gm = proxytensor_trace(mod, inputs)
        inputs_new = [torch.ones(5) + 5, torch.ones(5) + 8]
        output = gm(*inputs_new)
        ref_output = mod(*inputs_new)
        self.assertEqual(output, ref_output)
//  [python]
def test_leaf_operator_reg(self):
        class Leaf(torch.nn.Module):
            def forward(self, x, y):
                return x + y + torch.nn.Parameter(torch.ones(5))

        leaf = Leaf()
        wrap_lib.define("wrapped_foo(Tensor x, Tensor y) -> Tensor")
        wrap_lib.impl("wrapped_foo", leaf, "CPU")

        class Bar(torch.nn.Module):
            def __init__(self):
                super(Bar, self).__init__()
                self.foo = torch.ops.wrap.wrapped_foo
                self.other = torch.nn.Parameter(torch.ones(5))

            def forward(self, x, y):
                x = self.foo(x, y)
                x = x + self.other
                return x

        mod = Bar().eval()
        inputs = [torch.ones(5), torch.ones(5)]
        gm = proxytensor_trace(mod, inputs)
        inputs_new = [torch.ones(5) + 5, torch.ones(5) + 8]
        output = gm(*inputs_new)
        ref_output = mod(*inputs_new)
        torch.testing.assert_close(output, ref_output)
//  [python]
def test_reference_copy(self):
        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x, y):
                y[:, 0] = x[:, 0]
                return y

        mod = TestModule()

        def f(x, y):
            return mod(x, y)

        a = torch.ones(2, 2) + 2
        b = torch.ones(2, 2)
        b_copy = torch.ones(2, 2)
        ref_output = f(a, b)
        gm = make_fx(functionalize(f))(a, b)
        output = gm(a, b_copy)
        self.assertEqual(output, ref_output)
//  [python]
def test_reference_copy(self):
        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x, y):
                y[:, 0] = x[:, 0]
                return y

        mod = TestModule()

        def f(x, y):
            return mod(x, y)

        a = torch.ones(2, 2) + 2
        b = torch.ones(2, 2)
        b_copy = torch.ones(2, 2)
        ref_output = f(a, b)
        gm = make_fx(functionalize(f))(a, b)
        output = gm(a, b_copy)
        torch.testing.assert_close(output, ref_output)
//  [python]
def test_simple(self):
        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.relu = torch.nn.ReLU(inplace=True)

            def forward(self, x, y):
                y = y + x
                y = y.mul(x)
                y = y + x
                y = y + x
                y = y / x
                y = y + x
                y = y + x
                y = y / x
                y = y + x
                y = self.relu(y)
                return y

        mod = TestModule()
        mod = mod.cuda().half().eval()

        def f(x, y):
            return mod(x, y)

        inputs = [torch.randn(2, 5), torch.ones(2, 5)]
        inputs = [i.cuda().half() for i in inputs]
        ref_output = f(*inputs)

        mod = compile(
            mod,
            inputs,
            max_batch_size=100,
            explicit_batch_dimension=True,
            lower_precision=LowerPrecision.FP16,
            verbose_log=False,
            timing_cache_prefix="",
            save_timing_cache=False,
            cuda_graph_batch_size=-1,
            dynamic_batch=True,
            is_aten=True,
        )
        output = mod(*inputs)
        torch.testing.assert_allclose(output, ref_output)
//  [python]
def test_simple(self):
        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.relu = torch.nn.ReLU(inplace=True)

            def forward(self, x, y):
                y = y + x
                y = y.mul(x)
                y = y + x
                y = y + x
                y = y / x
                y = y + x
                y = y + x
                y = y / x
                y = y + x
                y = self.relu(y)
                return y

        mod = TestModule()
        mod = mod.cuda().half().eval()

        def f(x, y):
            return mod(x, y)

        inputs = [torch.randn(2, 5), torch.ones(2, 5)]
        inputs = [i.cuda().half() for i in inputs]
        ref_output = f(*inputs)

        mod = compile(
            mod,
            inputs,
            max_batch_size=100,
            explicit_batch_dimension=True,
            lower_precision=LowerPrecision.FP16,
            verbose_log=False,
            timing_cache_prefix="",
            save_timing_cache=False,
            cuda_graph_batch_size=-1,
            dynamic_batch=True,
            is_aten=True,
        )
        output = mod(*inputs)
        torch.testing.assert_close(output, ref_output)
//  [python]
def test_amp_fused_optimizer_parity(tmp_path):
    def run(fused=False):
        seed_everything(1234)
        model = FusedOptimizerParityModel(fused)
        trainer = Trainer(
            default_root_dir=tmp_path,
            accelerator="cuda",
            devices=1,
            precision="16-mixed",
            max_steps=5,
            logger=False,
            enable_checkpointing=False,
            enable_progress_bar=False,
            enable_model_summary=False,
        )
        trainer.fit(model)
        return model.parameters()

    params = run(fused=False)
    params_fused = run(fused=True)

    # Both the regular and the fused version of Adam produce the same losses and model weights
    for p, q in zip(params, params_fused):
        torch.testing.assert_allclose(p, q)
//  [python]
def test_amp_fused_optimizer_parity(tmp_path):
    def run(fused=False):
        seed_everything(1234)
        model = FusedOptimizerParityModel(fused)
        trainer = Trainer(
            default_root_dir=tmp_path,
            accelerator="cuda",
            devices=1,
            precision="16-mixed",
            max_steps=5,
            logger=False,
            enable_checkpointing=False,
            enable_progress_bar=False,
            enable_model_summary=False,
        )
        trainer.fit(model)
        return model.parameters()

    params = run(fused=False)
    params_fused = run(fused=True)

    # Both the regular and the fused version of Adam produce the same losses and model weights
    for p, q in zip(params, params_fused):
        torch.testing.assert_close(p, q)
//  [python]
def test_leaf_operator_reg(self):
        class Leaf(torch.nn.Module):
            def forward(self, x, y):
                return x + y + torch.nn.Parameter(torch.ones(5))

        leaf = Leaf()
        wrap_lib.define("wrapped_leaf(Tensor x, Tensor y) -> Tensor")
        wrap_lib.impl("wrapped_leaf", leaf, "CPU")

        class Bar(torch.nn.Module):
            def __init__(self):
                super(Bar, self).__init__()
                self.leaf = torch.ops.wrap.wrapped_leaf
                self.other = torch.nn.Parameter(torch.ones(5))

            def forward(self, x, y):
                x = self.leaf(x, y)
                x = x + self.other
                return x

        mod = Bar()

        def f(x, y):
            return mod(x, y)

        gm = make_fx(functionalize(f))(torch.ones(5), torch.ones(5))
        inputs = [torch.ones(5) + 5, torch.ones(5) + 8]
        output = gm(*inputs)
        ref_output = f(*inputs)
        t.assert_close(output, ref_output)
        # through the op registration method, the module is defined in a call_function
        call_function_node = None
        for node in gm.graph.nodes:
            if (
                node.op == "call_function"
                and node.target == torch.ops.wrap.wrapped_leaf
            ):
                call_function_node = node
        self.assertIsNotNone(call_function_node)
//  [python]
def test_leaf_operator_reg(self):
        class Leaf(torch.nn.Module):
            def forward(self, x, y):
                return x + y + torch.nn.Parameter(torch.ones(5))

        leaf = Leaf()
        wrap_lib.define("wrapped_leaf(Tensor x, Tensor y) -> Tensor")
        wrap_lib.impl("wrapped_leaf", leaf, "CPU")

        class Bar(torch.nn.Module):
            def __init__(self):
                super(Bar, self).__init__()
                self.leaf = torch.ops.wrap.wrapped_leaf
                self.other = torch.nn.Parameter(torch.ones(5))

            def forward(self, x, y):
                x = self.leaf(x, y)
                x = x + self.other
                return x

        mod = Bar()

        def f(x, y):
            return mod(x, y)

        gm = make_fx(functionalize(f))(torch.ones(5), torch.ones(5))
        inputs = [torch.ones(5) + 5, torch.ones(5) + 8]
        output = gm(*inputs)
        ref_output = f(*inputs)
        torch.testing.assert_close(output, ref_output)
        # through the op registration method, the module is defined in a call_function
        call_function_node = None
        for node in gm.graph.nodes:
            if (
                node.op == "call_function"
                and node.target == torch.ops.wrap.wrapped_leaf
            ):
                call_function_node = node
        self.assertIsNotNone(call_function_node)
//  [python]
def test_reference_copy_torchdynamo(self):
        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.relu = torch.nn.ReLU(inplace=True)

            def forward(self, x, y):
                y = y + 3
                y = self.relu(y)
                y[:, 0] = x[:, 0]
                return y

        mod = TestModule()

        def f(x, y):
            return mod(x, y)

        a = torch.ones(2, 2) + 2
        b = torch.ones(2, 2)
        inputs = [a, b]
        ref_output = f(*inputs)

        def compile_dispatch(gm, example_inputs):
            # dispatch tracer
            nargs = len(example_inputs)

            def fake_signature(fn, nargs):
                """FX gets confused by varargs, de-confuse it"""
                argnames = ",".join(f"arg{i}" for i in range(nargs))
                return eval(f"lambda {argnames}: fn({argnames})", {"fn": fn})

            gm = make_fx(functionalize(fake_signature(gm, nargs)))(*example_inputs)
            return gm

        optimized_mod = torchdynamo.optimize(
            compile_dispatch,
            nopython=True,
        )(mod)
        output = optimized_mod(*inputs)
        t.assert_close(output, ref_output)
//  [python]
def test_reference_copy_torchdynamo(self):
        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.relu = torch.nn.ReLU(inplace=True)

            def forward(self, x, y):
                y = y + 3
                y = self.relu(y)
                y[:, 0] = x[:, 0]
                return y

        mod = TestModule()

        def f(x, y):
            return mod(x, y)

        a = torch.ones(2, 2) + 2
        b = torch.ones(2, 2)
        inputs = [a, b]
        ref_output = f(*inputs)

        def compile_dispatch(gm, example_inputs):
            # dispatch tracer
            nargs = len(example_inputs)

            def fake_signature(fn, nargs):
                """FX gets confused by varargs, de-confuse it"""
                argnames = ",".join(f"arg{i}" for i in range(nargs))
                return eval(f"lambda {argnames}: fn({argnames})", {"fn": fn})

            gm = make_fx(functionalize(fake_signature(gm, nargs)))(*example_inputs)
            return gm

        optimized_mod = torchdynamo.optimize(
            compile_dispatch,
            nopython=True,
        )(mod)
        output = optimized_mod(*inputs)
        torch.testing.assert_close(output, ref_output)
//  [python]
def test_channel_shuffle_return_alias_of_self(self):
        # gh-76616: nn.ChannelShuffle will return alias of self with an empty input tensor
        groups = 3
        input_tensor = torch.rand([0, 9, 4, 4])
        output = torch.nn.ChannelShuffle(groups)(input_tensor)
        torch.testing.assert_allclose(output, input_tensor)
//  [python]
def test_channel_shuffle_return_alias_of_self(self):
        # gh-76616: nn.ChannelShuffle will return alias of self with an empty input tensor
        groups = 3
        input_tensor = torch.rand([0, 9, 4, 4])
        output = torch.nn.ChannelShuffle(groups)(input_tensor)
        torch.testing.assert_close(output, input_tensor)
//  [python]
def test_native_channel_shuffle_return_alias_of_self(self):
        groups = 3
        input_tensor = torch.rand([0, 9, 4, 4])
        output = torch.native_channel_shuffle(input_tensor, groups)
        self.assertEqual(output, input_tensor)
//  [python]
def test_native_channel_shuffle_return_alias_of_self(self):
        groups = 3
        input_tensor = torch.rand([0, 9, 4, 4])
        output = torch.native_channel_shuffle(input_tensor, groups)
        torch.testing.assert_close(output, input_tensor)
//  [python]
def test_pad_zeros(self):
        emb = Embeddings(
            embedding_dim=self.emb_size,
            vocab_size=self.vocab_size,
            padding_idx=self.pad_idx,
        )
        # pad embedding should be zeros
        torch.testing.assert_allclose(
            emb.lut.weight[self.pad_idx], torch.zeros([self.emb_size])
        )
//  [python]
def test_pad_zeros(self):
        emb = Embeddings(
            embedding_dim=self.emb_size,
            vocab_size=self.vocab_size,
            padding_idx=self.pad_idx,
        )
        # pad embedding should be zeros
        torch.testing.assert_close(
            emb.lut.weight[self.pad_idx], torch.zeros([self.emb_size])
        )
//  [python]
def test_reward_norm(self):
        torch.manual_seed(0)
        trainer = mocking_trainer()

        reward_normalizer = RewardNormalizer()
        reward_normalizer.register(trainer)

        batch = 10
        reward = torch.randn(batch, 1)
        td = TensorDict({REWARD_KEY: reward.clone()}, [batch])
        td_out = trainer._process_batch_hook(td)
        assert (td_out.get(REWARD_KEY) == reward).all()
        assert not reward_normalizer._normalize_has_been_called

        td_norm = trainer._process_optim_batch_hook(td)
        assert reward_normalizer._normalize_has_been_called
        torch.testing.assert_allclose(td_norm.get(REWARD_KEY).mean(), torch.zeros([]))
        torch.testing.assert_allclose(td_norm.get(REWARD_KEY).std(), torch.ones([]))
//  [python]
def test_reward_norm(self):
        torch.manual_seed(0)
        trainer = mocking_trainer()

        reward_normalizer = RewardNormalizer()
        reward_normalizer.register(trainer)

        batch = 10
        reward = torch.randn(batch, 1)
        td = TensorDict({REWARD_KEY: reward.clone()}, [batch])
        td_out = trainer._process_batch_hook(td)
        assert (td_out.get(REWARD_KEY) == reward).all()
        assert not reward_normalizer._normalize_has_been_called

        td_norm = trainer._process_optim_batch_hook(td)
        assert reward_normalizer._normalize_has_been_called
        torch.testing.assert_close(td_norm.get(REWARD_KEY).mean(), torch.zeros([]))
        torch.testing.assert_close(td_norm.get(REWARD_KEY).std(), torch.ones([]))
//  [python]
def test_forward(self):
        # fix the embedding weights
        weights = self._get_random_embedding_weights()
        emb = Embeddings(
            embedding_dim=self.emb_size,
            vocab_size=self.vocab_size,
            padding_idx=self.pad_idx,
        )
        emb.lut.weight.data = weights  # fill embeddings
        indices = torch.Tensor([0, 1, self.pad_idx, 9]).long()
        embedded = emb.forward(x=indices)
        # embedding operation is just slicing from weights matrix
        torch.testing.assert_allclose(
            embedded, torch.index_select(input=weights, index=indices, dim=0)
        )
        # after embedding, representations for PAD should still be zero
        torch.testing.assert_allclose(embedded[2], torch.zeros([self.emb_size]))
//  [python]
def test_forward(self):
        # fix the embedding weights
        weights = self._get_random_embedding_weights()
        emb = Embeddings(
            embedding_dim=self.emb_size,
            vocab_size=self.vocab_size,
            padding_idx=self.pad_idx,
        )
        emb.lut.weight.data = weights  # fill embeddings
        indices = torch.Tensor([0, 1, self.pad_idx, 9]).long()
        embedded = emb.forward(x=indices)
        # embedding operation is just slicing from weights matrix
        torch.testing.assert_close(
            embedded, torch.index_select(input=weights, index=indices, dim=0)
        )
        # after embedding, representations for PAD should still be zero
        torch.testing.assert_close(embedded[2], torch.zeros([self.emb_size]))
//  [python]
def test_scale(self):
        # fix the embedding weights
        weights = self._get_random_embedding_weights()
        emb = Embeddings(
            embedding_dim=self.emb_size,
            vocab_size=self.vocab_size,
            padding_idx=self.pad_idx,
            scale=True,
        )
        emb.lut.weight.data = weights
        indices = torch.Tensor([0, 1, self.pad_idx, 9]).long()
        embedded = emb.forward(x=indices)

        # now scaled
        # yapf: disable
        expected = torch.index_select(input=weights, index=indices, dim=0)
        self.assertEqual(embedded, expected * (self.emb_size**0.5))
//  [python]
def test_scale(self):
        # fix the embedding weights
        weights = self._get_random_embedding_weights()
        emb = Embeddings(
            embedding_dim=self.emb_size,
            vocab_size=self.vocab_size,
            padding_idx=self.pad_idx,
            scale=True,
        )
        emb.lut.weight.data = weights
        indices = torch.Tensor([0, 1, self.pad_idx, 9]).long()
        embedded = emb.forward(x=indices)

        # now scaled
        # yapf: disable
        expected = torch.index_select(input=weights, index=indices, dim=0)
        torch.testing.assert_close(
            embedded, expected * (self.emb_size**0.5), rtol=1e-4, atol=1e-4)
//  [python]
def assert_all_approx_close(a, b, atol=1e-8, rtol=1e-5, count=10):
    idx = torch.isclose(a, b, rtol=rtol, atol=atol)
    sumval = (idx == 0).sum().item()
    if sumval > count:
        print(f"Too many values not close: assert {sumval} < {count}")
        torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
//  [python]
def test_tanhtrsf_grad(dtype):
    torch.manual_seed(0)
    trsf = SafeTanhTransform()
    x = torch.randn(100, requires_grad=True)
    y1 = trsf(x)
    y2 = x.tanh()
    g1 = autograd.grad(y1.sum(), x, retain_graph=True)[0]
    g2 = autograd.grad(y2.sum(), x, retain_graph=True)[0]
    torch.testing.assert_allclose(g1, g2)
//  [python]
def test_tanhtrsf_grad(dtype):
    torch.manual_seed(0)
    trsf = SafeTanhTransform()
    x = torch.randn(100, requires_grad=True)
    y1 = trsf(x)
    y2 = x.tanh()
    g1 = autograd.grad(y1.sum(), x, retain_graph=True)[0]
    g2 = autograd.grad(y2.sum(), x, retain_graph=True)[0]
    torch.testing.assert_close(g1, g2)
//  [python]
def test_sample(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedCategorical(
            probs=probs,
            mask=mask,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples])
        sample_probs = torch.bincount(samples) / num_samples
        self.assertEqual(sample_probs.size(), ref_probs.size())
        diff = (sample_probs - ref_probs).abs()
        self.assertLess(diff.max(), 1e-2)
//  [python]
def test_sample(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedCategorical(
            probs=probs,
            mask=mask,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples])
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_close(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_sample_sparse(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        indices = torch.tensor([0, 2, 3])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedCategorical(
            logits=logits,
            indices=indices,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples])
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_allclose(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_sample_sparse(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        indices = torch.tensor([0, 2, 3])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedCategorical(
            logits=logits,
            indices=indices,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples])
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_close(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_amp_fused_optimizer_parity():
    def run(fused=False):
        seed_everything(1234)
        fabric = Fabric(accelerator="cuda", precision=16, devices=1)

        model = nn.Linear(10, 10).to(fabric.device)  # TODO: replace with individual setup_module call
        optimizer = torch.optim.Adam(model.parameters(), lr=1.0, fused=fused)

        model, optimizer = fabric.setup(model, optimizer)
        assert isinstance(fabric._precision.scaler, torch.cuda.amp.GradScaler)

        data = torch.randn(10, 10, device="cuda")
        target = torch.randn(10, 10, device="cuda")

        losses = []
        for _ in range(5):
            optimizer.zero_grad()
            output = model(data)
            loss = (output - target).abs().sum()
            fabric.backward(loss)
            optimizer.step()
            losses.append(loss.detach())
        return torch.stack(losses), model.parameters()

    losses, params = run(fused=False)
    losses_fused, params_fused = run(fused=True)

    # Both the regular and the fused version of Adam produce the same losses and model weights
    torch.testing.assert_allclose(losses, losses_fused)
    for p, q in zip(params, params_fused):
        torch.testing.assert_allclose(p, q)
//  [python]
def test_amp_fused_optimizer_parity():
    def run(fused=False):
        seed_everything(1234)
        fabric = Fabric(accelerator="cuda", precision=16, devices=1)

        model = nn.Linear(10, 10).to(fabric.device)  # TODO: replace with individual setup_module call
        optimizer = torch.optim.Adam(model.parameters(), lr=1.0, fused=fused)

        model, optimizer = fabric.setup(model, optimizer)
        assert isinstance(fabric._precision.scaler, torch.cuda.amp.GradScaler)

        data = torch.randn(10, 10, device="cuda")
        target = torch.randn(10, 10, device="cuda")

        losses = []
        for _ in range(5):
            optimizer.zero_grad()
            output = model(data)
            loss = (output - target).abs().sum()
            fabric.backward(loss)
            optimizer.step()
            losses.append(loss.detach())
        return torch.stack(losses), model.parameters()

    losses, params = run(fused=False)
    losses_fused, params_fused = run(fused=True)

    # Both the regular and the fused version of Adam produce the same losses and model weights
    torch.testing.assert_close(losses, losses_fused)
    for p, q in zip(params, params_fused):
        torch.testing.assert_close(p, q)
//  [python]
def test_one_hot(self):
        torch.manual_seed(0)
        logits = torch.randn(1, 10)
        torch.manual_seed(0)
        d = OneHotCategorical(logits=logits)
        s_a = d.sample((1,))
        torch.manual_seed(0)
        d = OneHotCategorical(probs=torch.softmax(logits, -1))
        s_b = d.sample((1,))
        torch.testing.assert_allclose(s_a, s_b)
        assert s_a.dtype == torch.long
        assert s_b.dtype == torch.long
        assert s_a.sum(-1) == 1
        assert s_b.sum(-1) == 1
        assert s_a.shape[-1] == 10
        assert s_b.shape[-1] == 10
//  [python]
def test_one_hot(self):
        torch.manual_seed(0)
        logits = torch.randn(1, 10)
        torch.manual_seed(0)
        d = OneHotCategorical(logits=logits)
        s_a = d.sample((1,))
        torch.manual_seed(0)
        d = OneHotCategorical(probs=torch.softmax(logits, -1))
        s_b = d.sample((1,))
        torch.testing.assert_close(s_a, s_b)
        assert s_a.dtype == torch.long
        assert s_b.dtype == torch.long
        assert s_a.sum(-1) == 1
        assert s_b.sum(-1) == 1
        assert s_a.shape[-1] == 10
        assert s_b.shape[-1] == 10
//  [python]
def test_sample(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedOneHotCategorical(
            probs=probs,
            mask=mask,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples]).argmax(-1)
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_allclose(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_sample(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedOneHotCategorical(
            probs=probs,
            mask=mask,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples]).argmax(-1)
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_close(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_sample_sparse(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        indices = torch.tensor([0, 2, 3])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedOneHotCategorical(
            logits=logits,
            indices=indices,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples]).argmax(-1)
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_allclose(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_sample_sparse(self, neg_inf: float) -> None:
        torch.manual_seed(0)
        logits = torch.randn(4)
        probs = F.softmax(logits, dim=-1)
        mask = torch.tensor([True, False, True, True])
        indices = torch.tensor([0, 2, 3])
        ref_probs = probs.masked_fill(~mask, 0.0)
        ref_probs /= ref_probs.sum(dim=-1, keepdim=True)

        dist = MaskedOneHotCategorical(
            logits=logits,
            indices=indices,
            neg_inf=neg_inf,
        )
        num_samples = 10000
        samples = dist.sample([num_samples]).argmax(-1)
        sample_probs = torch.bincount(samples) / num_samples
        torch.testing.assert_close(sample_probs, ref_probs, rtol=1e-5, atol=1e-2)
//  [python]
def test_init():
    base_logits = torch.Tensor([0.3, 0.2, 0.1])

    # no batch
    x_hat = torch.Tensor([0.3, 0.2, 0.1])
    noise_model = BernoulliNoise(x_hat, base_logits)

    assert noise_model.logits.shape == torch.Size([3])
    assert torch.equal(noise_model.logits, torch.Tensor([0.6, 0.4, 0.2]))

    # batch size 1
    x_hat = torch.Tensor([[0.3, 0.2, 0.1], [0.4, 0.5, 0.6]])
    noise_model = BernoulliNoise(x_hat, base_logits)

    assert noise_model.logits.shape == torch.Size([2, 3])
    torch.testing.assert_allclose(noise_model.logits, torch.Tensor([[0.6, 0.4, 0.2], [0.7, 0.7, 0.7]]))
//  [python]
def test_init():
    base_logits = torch.Tensor([0.3, 0.2, 0.1])

    # no batch
    x_hat = torch.Tensor([0.3, 0.2, 0.1])
    noise_model = BernoulliNoise(x_hat, base_logits)

    assert noise_model.logits.shape == torch.Size([3])
    assert torch.equal(noise_model.logits, torch.Tensor([0.6, 0.4, 0.2]))

    # batch size 1
    x_hat = torch.Tensor([[0.3, 0.2, 0.1], [0.4, 0.5, 0.6]])
    noise_model = BernoulliNoise(x_hat, base_logits)

    assert noise_model.logits.shape == torch.Size([2, 3])
    torch.testing.assert_close(noise_model.logits, torch.Tensor([[0.6, 0.4, 0.2], [0.7, 0.7, 0.7]]))
//  [python]
def test_joint_noise_passthrough(noise: Noise):
    """Test that properties from the underlying noise distributions are passed through."""
    joint_noise = JointNoise({"a": noise})

    # Distribution properties
    t.assert_close(joint_noise.entropy(), noise.entropy())
    t.assert_close(joint_noise.mode.get("a"), noise.mode)
    t.assert_close(joint_noise.mean.get("a"), noise.mean)
//  [python]
def test_joint_noise_passthrough(noise: Noise):
    """Test that properties from the underlying noise distributions are passed through."""
    joint_noise = JointNoise({"a": noise})

    # Distribution properties
    torch.testing.assert_close(joint_noise.entropy(), noise.entropy())
    torch.testing.assert_close(joint_noise.mode.get("a"), noise.mode)
    torch.testing.assert_close(joint_noise.mean.get("a"), noise.mean)
//  [python]
def test_construction(self, neg_inf, sparse, logits, device):
        torch.manual_seed(0)
        logits_vals = torch.randn(4, device=device) / 100  # almost equal probabilities
        if logits:
            logits = logits_vals
            probs = None
        else:
            probs = logits_vals.softmax(-1)
            logits = None

        if sparse:
            indices = torch.tensor([0, 2, 3], device=device)
            mask = None
        else:
            mask = torch.tensor([True, False, True, True], device=device)
            indices = None
        dist = MaskedOneHotCategorical(
            logits=logits, probs=probs, indices=indices, mask=mask, neg_inf=neg_inf
        )
        dist_categ = MaskedCategorical(
            logits=logits, probs=probs, indices=indices, mask=mask, neg_inf=neg_inf
        )
        for _ in range(10):
            sample = dist.sample((100,))
            assert not sample[..., 1].any()
            assert torch.isfinite(dist.log_prob(sample)).all()
            torch.testing.assert_allclose(
                dist.log_prob(sample), dist_categ.log_prob(sample.argmax(-1))
            )
            assert sample.device == device

        sample_unfeasible = torch.zeros_like(sample)
        sample_unfeasible[..., 1] = 1
        if neg_inf == -float("inf"):
            assert (dist.log_prob(sample_unfeasible) == neg_inf).all()
        else:
            assert (dist.log_prob(sample_unfeasible) > -float("inf")).all()
//  [python]
def test_construction(self, neg_inf, sparse, logits, device):
        torch.manual_seed(0)
        logits_vals = torch.randn(4, device=device) / 100  # almost equal probabilities
        if logits:
            logits = logits_vals
            probs = None
        else:
            probs = logits_vals.softmax(-1)
            logits = None

        if sparse:
            indices = torch.tensor([0, 2, 3], device=device)
            mask = None
        else:
            mask = torch.tensor([True, False, True, True], device=device)
            indices = None
        dist = MaskedOneHotCategorical(
            logits=logits, probs=probs, indices=indices, mask=mask, neg_inf=neg_inf
        )
        dist_categ = MaskedCategorical(
            logits=logits, probs=probs, indices=indices, mask=mask, neg_inf=neg_inf
        )
        for _ in range(10):
            sample = dist.sample((100,))
            assert not sample[..., 1].any()
            assert torch.isfinite(dist.log_prob(sample)).all()
            torch.testing.assert_close(
                dist.log_prob(sample), dist_categ.log_prob(sample.argmax(-1))
            )
            assert sample.device == device

        sample_unfeasible = torch.zeros_like(sample)
        sample_unfeasible[..., 1] = 1
        if neg_inf == -float("inf"):
            assert (dist.log_prob(sample_unfeasible) == neg_inf).all()
        else:
            assert (dist.log_prob(sample_unfeasible) > -float("inf")).all()
//  [python]
def test_fill_kv_cache(self, k_states, v_states, k_caches, v_caches,
                           block_offsets, q_start_loc, q_seq_length,
                           kv_seq_length, max_q_seq_length, gt):
        fill_kv_cache(k_states, v_states, k_caches, v_caches, q_start_loc,
                      q_seq_length, kv_seq_length, max_q_seq_length,
                      block_offsets)

        self.assertEqual(k_caches, gt[0])
        self.assertEqual(v_caches, gt[1])
//  [python]
def test_fill_kv_cache(self, k_states, v_states, k_caches, v_caches,
                           block_offsets, q_start_loc, q_seq_length,
                           kv_seq_length, max_q_seq_length, gt):
        fill_kv_cache(k_states, v_states, k_caches, v_caches, q_start_loc,
                      q_seq_length, kv_seq_length, max_q_seq_length,
                      block_offsets)

        torch.testing.assert_close(k_caches, gt[0])
        torch.testing.assert_close(v_caches, gt[1])
//  [python]
def test_matching(self):
        actual = torch.tensor(1.0)
        expected = actual.clone()

        torch.testing.assert_allclose(actual, expected)
//  [python]
def test_matching(self):
        actual = torch.tensor(1.0)
        expected = actual.clone()

        torch.testing.assert_close(actual, expected)
//  [python]
def test_unexpected_error_originate(self, _):
        actual = torch.tensor(1.0)
        expected = actual.clone()

        with self.assertRaisesRegex(RuntimeError, "unexpected exception"):
            self.assertEqual(actual, expected)
//  [python]
def test_unexpected_error_originate(self, _):
        actual = torch.tensor(1.0)
        expected = actual.clone()

        with self.assertRaisesRegex(RuntimeError, "unexpected exception"):
            torch.testing.assert_close(actual, expected)
//  [python]
def test_transformed_env_constructor_with_state_dict(from_pixels):
    config_fields = [
        (config_field.name, config_field.type, config_field)
        for config_cls in (
            EnvConfig,
            DreamerConfig,
        )
        for config_field in dataclasses.fields(config_cls)
    ]
    flags = list(from_pixels)

    Config = dataclasses.make_dataclass(cls_name="Config", fields=config_fields)
    cs = ConfigStore.instance()
    cs.store(name="config", node=Config)
    with initialize(version_base="1.1", config_path=None):
        cfg = compose(config_name="config", overrides=flags)
        env_maker = (
            ContinuousActionConvMockEnvNumpy
            if from_pixels
            else ContinuousActionVecMockEnv
        )
        t_env = transformed_env_constructor(
            cfg,
            use_env_creator=False,
            custom_env_maker=env_maker,
        )()
        for t in t_env.transform:
            if isinstance(t, ObservationNorm):
                t.init_stats(4)
        idx, state_dict = retrieve_observation_norms_state_dict(t_env)[0]

        obs_transform = transformed_env_constructor(
            cfg,
            use_env_creator=False,
            custom_env_maker=env_maker,
            obs_norm_state_dict=state_dict,
        )().transform[idx]
        t.assert_close(obs_transform.state_dict(), state_dict)
//  [python]
def test_transformed_env_constructor_with_state_dict(from_pixels):
    config_fields = [
        (config_field.name, config_field.type, config_field)
        for config_cls in (
            EnvConfig,
            DreamerConfig,
        )
        for config_field in dataclasses.fields(config_cls)
    ]
    flags = list(from_pixels)

    Config = dataclasses.make_dataclass(cls_name="Config", fields=config_fields)
    cs = ConfigStore.instance()
    cs.store(name="config", node=Config)
    with initialize(version_base="1.1", config_path=None):
        cfg = compose(config_name="config", overrides=flags)
        env_maker = (
            ContinuousActionConvMockEnvNumpy
            if from_pixels
            else ContinuousActionVecMockEnv
        )
        t_env = transformed_env_constructor(
            cfg,
            use_env_creator=False,
            custom_env_maker=env_maker,
        )()
        for t in t_env.transform:
            if isinstance(t, ObservationNorm):
                t.init_stats(4)
        idx, state_dict = retrieve_observation_norms_state_dict(t_env)[0]

        obs_transform = transformed_env_constructor(
            cfg,
            use_env_creator=False,
            custom_env_maker=env_maker,
            obs_norm_state_dict=state_dict,
        )().transform[idx]
        torch.testing.assert_close(obs_transform.state_dict(), state_dict)
//  [python]
def test_joint_noise_properties(noise_a: Noise, noise_b: Noise):
    """Test that properties from the underlying noise distributions are passed through for pairs of noise."""
    joint_noise = JointNoise({"a": noise_a, "b": noise_b})

    # Distribution properties
    t.assert_close(joint_noise.entropy(), noise_a.entropy() + noise_b.entropy())
    t.assert_close(joint_noise.mode.get("a"), noise_a.mode)
    t.assert_close(joint_noise.mode.get("b"), noise_b.mode)
    t.assert_close(joint_noise.mean.get("a"), noise_a.mean)
    t.assert_close(joint_noise.mean.get("b"), noise_b.mean)
//  [python]
def test_joint_noise_properties(noise_a: Noise, noise_b: Noise):
    """Test that properties from the underlying noise distributions are passed through for pairs of noise."""
    joint_noise = JointNoise({"a": noise_a, "b": noise_b})

    # Distribution properties
    torch.testing.assert_close(joint_noise.entropy(), noise_a.entropy() + noise_b.entropy())
    torch.testing.assert_close(joint_noise.mode.get("a"), noise_a.mode)
    torch.testing.assert_close(joint_noise.mode.get("b"), noise_b.mode)
    torch.testing.assert_close(joint_noise.mean.get("a"), noise_a.mean)
    torch.testing.assert_close(joint_noise.mean.get("b"), noise_b.mean)
//  [python]
def test_unexpected_error_compare(self, _):
        actual = torch.tensor(1.0)
        expected = actual.clone()

        with self.assertRaisesRegex(RuntimeError, "unexpected exception"):
            self.assertEqual(actual, expected)
//  [python]
def test_unexpected_error_compare(self, _):
        actual = torch.tensor(1.0)
        expected = actual.clone()

        with self.assertRaisesRegex(RuntimeError, "unexpected exception"):
            torch.testing.assert_close(actual, expected)
//  [python]
def test_joint_noise_sample_log_prob(noise_a: Noise, noise_b: Noise, sample_shape: torch.Size):
    """Test that sampling and log prob work as expected for pairs of noise."""
    joint_noise = JointNoise({"a": noise_a, "b": noise_b})
    sample_a = noise_a.sample(sample_shape)
    sample_b = noise_b.sample(sample_shape)
    joint_sample = TensorDict({"a": sample_a, "b": sample_b}, batch_size=sample_shape)
    t.assert_close(
        joint_noise.log_prob(joint_sample), noise_a.log_prob(sample_a) + noise_b.log_prob(sample_b)
    )
    joint_sample = joint_noise.sample()
    sample_a = joint_sample.get("a")
    sample_b = joint_sample.get("b")
    t.assert_close(
        joint_noise.log_prob(joint_sample), noise_a.log_prob(sample_a) + noise_b.log_prob(sample_b)
    )
//  [python]
def test_joint_noise_sample_log_prob(noise_a: Noise, noise_b: Noise, sample_shape: torch.Size):
    """Test that sampling and log prob work as expected for pairs of noise."""
    joint_noise = JointNoise({"a": noise_a, "b": noise_b})
    sample_a = noise_a.sample(sample_shape)
    sample_b = noise_b.sample(sample_shape)
    joint_sample = TensorDict({"a": sample_a, "b": sample_b}, batch_size=sample_shape)
    torch.testing.assert_close(
        joint_noise.log_prob(joint_sample), noise_a.log_prob(sample_a) + noise_b.log_prob(sample_b)
    )
    joint_sample = joint_noise.sample()
    sample_a = joint_sample.get("a")
    sample_b = joint_sample.get("b")
    torch.testing.assert_close(
        joint_noise.log_prob(joint_sample), noise_a.log_prob(sample_a) + noise_b.log_prob(sample_b)
    )
//  [python]
def test_fp16(self, batch=4):
        groups = 32
        size = (12, 12)
        target = detect_target()
        X = Tensor(
            shape=[IntImm(batch), *size, 32],
            dtype="float16",
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[32, 3, 3, 1], dtype="float16", name="input_1", is_input=True)
        OP = ops.conv2d_depthwise(stride=1, pad=1, dilate=1, group=groups)
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(Y, target, "./tmp", "conv2d_dw")

        X_pt = torch.randn(batch, 32, *size).cuda().half()
        W_pt = torch.randn(32, 1, 3, 3).cuda().half()
        Y_pt = torch.nn.functional.conv2d(X_pt, W_pt, padding=1, groups=groups)
        x = X_pt.permute((0, 2, 3, 1)).contiguous()
        w = W_pt.permute((0, 2, 3, 1)).contiguous()
        y = torch.empty([batch, *size, 32]).cuda().half()
        module.run_with_tensors({"input_0": x, "input_1": w}, [y])
        y_transpose = y.permute((0, 3, 1, 2))
        self.assertFalse(y_transpose.isnan().any())
        self.assertFalse(y_transpose.isinf().any())
        if target.name() == "cuda":
            torch.testing.assert_allclose(Y_pt, y_transpose, atol=1e-2, rtol=1e-2)
        else:
            torch.testing.assert_allclose(Y_pt, y_transpose, atol=1.25e-1, rtol=1e-1)
//  [python]
def test_fp16(self, batch=4):
        groups = 32
        size = (12, 12)
        target = detect_target()
        X = Tensor(
            shape=[IntImm(batch), *size, 32],
            dtype="float16",
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[32, 3, 3, 1], dtype="float16", name="input_1", is_input=True)
        OP = ops.conv2d_depthwise(stride=1, pad=1, dilate=1, group=groups)
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(Y, target, "./tmp", "conv2d_dw")

        X_pt = torch.randn(batch, 32, *size).cuda().half()
        W_pt = torch.randn(32, 1, 3, 3).cuda().half()
        Y_pt = torch.nn.functional.conv2d(X_pt, W_pt, padding=1, groups=groups)
        x = X_pt.permute((0, 2, 3, 1)).contiguous()
        w = W_pt.permute((0, 2, 3, 1)).contiguous()
        y = torch.empty([batch, *size, 32]).cuda().half()
        module.run_with_tensors({"input_0": x, "input_1": w}, [y])
        y_transpose = y.permute((0, 3, 1, 2))
        self.assertFalse(y_transpose.isnan().any())
        self.assertFalse(y_transpose.isinf().any())
        if target.name() == "cuda":
            torch.testing.assert_close(Y_pt, y_transpose, atol=1e-2, rtol=1e-2)
        else:
            torch.testing.assert_close(Y_pt, y_transpose, atol=1.25e-1, rtol=1e-1)
//  [python]
def test_sequence_mismatching_len(self):
        actual = (torch.empty(()),)
        expected = ()

        with self.assertRaises(AssertionError):
            self.assertEqual(actual, expected)
//  [python]
def test_sequence_mismatching_len(self):
        actual = (torch.empty(()),)
        expected = ()

        with self.assertRaises(AssertionError):
            torch.testing.assert_close(actual, expected)
//  [python]
def test_mapping_mismatching_keys(self):
        actual = {"a": torch.empty(())}
        expected = {}

        with self.assertRaises(AssertionError):
            self.assertEqual(actual, expected)
//  [python]
def test_mapping_mismatching_keys(self):
        actual = {"a": torch.empty(())}
        expected = {}

        with self.assertRaises(AssertionError):
            torch.testing.assert_close(actual, expected)
//  [python]
def test_sequence_mismatching_values_msg(self):
        t1 = torch.tensor(1)
        t2 = torch.tensor(2)

        actual = (t1, t1)
        expected = (t1, t2)

        with self.assertRaisesRegex(AssertionError, re.escape("item [1]")):
            self.assertEqual(actual, expected)
//  [python]
def test_sequence_mismatching_values_msg(self):
        t1 = torch.tensor(1)
        t2 = torch.tensor(2)

        actual = (t1, t1)
        expected = (t1, t2)

        with self.assertRaisesRegex(AssertionError, re.escape("item [1]")):
            torch.testing.assert_close(actual, expected)
//  [python]
def test_mapping_mismatching_values_msg(self):
        t1 = torch.tensor(1)
        t2 = torch.tensor(2)

        actual = {"a": t1, "b": t1}
        expected = {"a": t1, "b": t2}

        with self.assertRaisesRegex(AssertionError, re.escape("item ['b']")):
            self.assertEqual(actual, expected)
//  [python]
def test_mapping_mismatching_values_msg(self):
        t1 = torch.tensor(1)
        t2 = torch.tensor(2)

        actual = {"a": t1, "b": t1}
        expected = {"a": t1, "b": t2}

        with self.assertRaisesRegex(AssertionError, re.escape("item ['b']")):
            torch.testing.assert_close(actual, expected)
//  [python]
def test_low_high_boolean_integral1(self, dtype, device):
        shape = (10_000,)
        eps = 1e-4

        actual = torch.testing.make_tensor(shape, dtype=dtype, device=device, low=-(1 - eps), high=1 - eps)
        expected = torch.zeros(shape, dtype=dtype, device=device)

        self.assertEqual(actual, expected)
//  [python]
def test_low_high_boolean_integral1(self, dtype, device):
        shape = (10_000,)
        eps = 1e-4

        actual = torch.testing.make_tensor(shape, dtype=dtype, device=device, low=-(1 - eps), high=1 - eps)
        expected = torch.zeros(shape, dtype=dtype, device=device)

        torch.testing.assert_close(actual, expected)
//  [python]
def _assert_weights_equivalent(self, m1: torch.nn.Module, m2: torch.nn.Module):
        for p1, p2 in zip(m1.parameters(), m2.parameters()):
            torch.testing.assert_allclose(p1, p2)
//  [python]
def _assert_weights_equivalent(self, m1: torch.nn.Module, m2: torch.nn.Module):
        for p1, p2 in zip(m1.parameters(), m2.parameters()):
            torch.testing.assert_close(p1, p2)