jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_split_module_default_arg(self):
        class ModelToTrace(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.lin = torch.nn.Linear(512, 512)

            def forward(self, x, targets=None):
                x = self.lin(x)

                if targets is not None:
                    x = x + targets

                return x

        mtt = ModelToTrace()
        traced = torch.fx.symbolic_trace(mtt, concrete_args={'targets': None})

        split = split_module(traced, mtt, lambda node: 0)

        x = torch.randn(50, 512)
        torch.testing.assert_close(split(x), traced(x))
//  [python]
def test_selective_checkpoint_wrapper_compile(
    device, no_grad, memory_budget, inplace, random
):
    torch.manual_seed(42)
    dtype = torch.float16
    modules = _get_model_blocks(
        3, dtype, device, inplace=inplace, random=random, first_inplace=False
    )
    inputs = torch.rand(32, 128, 10, dtype=dtype, device=device)

    model = torch.nn.Sequential(
        *[selective_checkpoint_wrapper(b, memory_budget=memory_budget) for b in modules]
    )
    model = torch.compile(model)
    model_ref = torch.nn.Sequential(*deepcopy(modules))

    grad = torch.rand_like(inputs)

    context = torch.no_grad() if no_grad else nullcontext()

    with context:
        torch.manual_seed(42)
        out = model(inputs.clone())
        if not no_grad:
            out.backward(grad)

        torch.manual_seed(42)
        out_ref = model_ref(inputs.clone())
        if not no_grad:
            out_ref.backward(grad)

    atol = 3e-4
    rtol = 1e-3
    torch.testing.assert_allclose(out, out_ref, atol=atol, rtol=rtol)

    if no_grad:
        return

    for p, p_ref in zip(model.parameters(), model_ref.parameters()):
        atol = 4e-4
        rtol = 2e-3
        torch.testing.assert_allclose(p.grad, p_ref.grad, atol=atol, rtol=rtol)
//  [python]
def test_selective_checkpoint_wrapper_compile(
    device, no_grad, memory_budget, inplace, random
):
    torch.manual_seed(42)
    dtype = torch.float16
    modules = _get_model_blocks(
        3, dtype, device, inplace=inplace, random=random, first_inplace=False
    )
    inputs = torch.rand(32, 128, 10, dtype=dtype, device=device)

    model = torch.nn.Sequential(
        *[selective_checkpoint_wrapper(b, memory_budget=memory_budget) for b in modules]
    )
    model = torch.compile(model)
    model_ref = torch.nn.Sequential(*deepcopy(modules))

    grad = torch.rand_like(inputs)

    context = torch.no_grad() if no_grad else nullcontext()

    with context:
        torch.manual_seed(42)
        out = model(inputs.clone())
        if not no_grad:
            out.backward(grad)

        torch.manual_seed(42)
        out_ref = model_ref(inputs.clone())
        if not no_grad:
            out_ref.backward(grad)

    atol = 3e-4
    rtol = 1e-3
    torch.testing.assert_close(out, out_ref, atol=atol, rtol=rtol)

    if no_grad:
        return

    for p, p_ref in zip(model.parameters(), model_ref.parameters()):
        atol = 4e-4
        rtol = 2e-3
        torch.testing.assert_close(p.grad, p_ref.grad, atol=atol, rtol=rtol)
//  [python]
def test_meta_tracer(self):
        class MetaTracerTestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.emb = torch.nn.Embedding(num_embeddings=42, embedding_dim=16)
                self.layernorm = torch.nn.LayerNorm(16)

            def forward(self, x):
                emb = self.emb(x)
                emb = emb + torch.arange(emb.shape[-1], dtype=torch.float, device=emb.device)
                lol = self.layernorm(emb)
                return torch.relu(lol) if lol.shape[0] < 30 else torch.sigmoid(lol)

        mttm = MetaTracerTestModule()
        for BS in [15, 35]:
            x = torch.zeros(BS, dtype=torch.long).random_(42)
            meta_args = {'x' : x.to(device='meta')}
            gm = torch.fx.experimental.meta_tracer.symbolic_trace(mttm, meta_args=meta_args)
            t.assert_close(gm(x), mttm(x))

            # Test serialization/deserialization
            with tempfile.TemporaryDirectory() as tmp_dir:
                with open(f'{tmp_dir}/meta_module.pkl', 'wb') as f:
                    pickle.dump(gm, f)

                with open(f'{tmp_dir}/meta_module.pkl', 'rb') as f:
                    loaded = pickle.load(f)

                t.assert_close(loaded(x), mttm(x))
//  [python]
def test_meta_tracer(self):
        class MetaTracerTestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.emb = torch.nn.Embedding(num_embeddings=42, embedding_dim=16)
                self.layernorm = torch.nn.LayerNorm(16)

            def forward(self, x):
                emb = self.emb(x)
                emb = emb + torch.arange(emb.shape[-1], dtype=torch.float, device=emb.device)
                lol = self.layernorm(emb)
                return torch.relu(lol) if lol.shape[0] < 30 else torch.sigmoid(lol)

        mttm = MetaTracerTestModule()
        for BS in [15, 35]:
            x = torch.zeros(BS, dtype=torch.long).random_(42)
            meta_args = {'x' : x.to(device='meta')}
            gm = torch.fx.experimental.meta_tracer.symbolic_trace(mttm, meta_args=meta_args)
            torch.testing.assert_close(gm(x), mttm(x))

            # Test serialization/deserialization
            with tempfile.TemporaryDirectory() as tmp_dir:
                with open(f'{tmp_dir}/meta_module.pkl', 'wb') as f:
                    pickle.dump(gm, f)

                with open(f'{tmp_dir}/meta_module.pkl', 'rb') as f:
                    loaded = pickle.load(f)

                torch.testing.assert_close(loaded(x), mttm(x))
//  [python]
def test_normalize_binary_operators(self):
        ops_to_test = {
            torch.add,
            torch.mul,
            torch.sub,
            torch.div,
            torch.floor_divide,
            torch.remainder,
            torch.eq,
            torch.ne,
            torch.lt,
            torch.le,
            torch.gt,
            torch.ge,
        }

        # Test Tensor/Tensor callsite
        for op in ops_to_test:

            class WrapperMod(torch.nn.Module):
                def forward(self, x, y):
                    return op(x, y)

            traced = symbolic_trace(WrapperMod())
            normalized = NormalizeOperators(traced).transform()
            x, y = torch.randn(3, 4), torch.randn(3, 4)
            self.assertEqual(traced(x, y), normalized(x, y))
            self.assertFalse(
                any(n.target in ops_to_test for n in normalized.graph.nodes)
            )

        # Test Tensor/scalar callsite
        for op in ops_to_test:

            class WrapperMod(torch.nn.Module):
                def forward(self, x):
                    return op(x, 42)

            traced = symbolic_trace(WrapperMod())
            normalized = NormalizeOperators(traced).transform()
            x = torch.randn(3, 4)
            self.assertEqual(traced(x), normalized(x))
            self.assertFalse(
                any(n.target in ops_to_test for n in normalized.graph.nodes)
            )
//  [python]
def test_optimize_for_inference_cpu(self):
        import torch.nn as nn

        class Foo(nn.Module):
            def __init__(self):
                super().__init__()
                layers = []
                layers2 = []
                for _ in range(10):
                    layers.append(nn.Conv2d(3, 3, 1))
                    layers.append(nn.BatchNorm2d(3))
                    layers.append(nn.ReLU())

                    layers2.append(nn.Conv2d(3, 3, 1))
                    layers2.append(nn.BatchNorm2d(3))
                    layers2.append(nn.ReLU())
                self.model = nn.Sequential(*layers)
                self.model2 = nn.Sequential(*layers2)

            def forward(self, x):
                return self.model(x) + self.model2(x)

        N, C, H, W, = (
            1,
            3,
            224,
            224,
        )
        inp = torch.randn(N, C, H, W)
        with torch.no_grad():
            model = Foo().eval()
            optimized_model = optimization.optimize_for_inference(model)
            t.assert_close(model(inp), optimized_model(inp))

            optimized_model2 = optimization.optimize_for_inference(
                model, pass_config={"remove_dropout": False}
            )
            t.assert_close(model(inp), optimized_model2(inp))
//  [python]
def test_gradient_magnitude():
    constant = torch.ones((5, 2, 4, 4))
    constant_grad_magnitude = fme.gradient_magnitude(constant, dim=(-2, -1))
    t.assert_close(constant_grad_magnitude, torch.zeros_like(constant))

    monotonic = torch.tile(torch.arange(4.0), (5, 2, 4, 1))
    monotonic_grad_magnitude = fme.gradient_magnitude(monotonic, dim=(-2, -1))
    t.assert_close(monotonic_grad_magnitude, torch.ones_like(monotonic))
//  [python]
def test_gradient_magnitude():
    constant = torch.ones((5, 2, 4, 4))
    constant_grad_magnitude = fme.gradient_magnitude(constant, dim=(-2, -1))
    torch.testing.assert_close(constant_grad_magnitude, torch.zeros_like(constant))

    monotonic = torch.tile(torch.arange(4.0), (5, 2, 4, 1))
    monotonic_grad_magnitude = fme.gradient_magnitude(monotonic, dim=(-2, -1))
    torch.testing.assert_close(monotonic_grad_magnitude, torch.ones_like(monotonic))
//  [python]
def test_rmse_of_time_mean(variable, time, lats, n_lon):
    x = torch.randn(variable, time, len(lats), n_lon)
    random_weights = torch.rand(len(lats), n_lon)

    result = fme.rmse_of_time_mean(x, x.clone(), time_dim=1)
    t.assert_close(
        result,
        torch.zeros((variable,)),
        msg="RMSE of time mean between identical tensors should be zero.",
    )

    result = fme.rmse_of_time_mean(
        torch.zeros(variable, time, len(lats), n_lon),
        torch.ones(variable, time, len(lats), n_lon),
        weights=random_weights,
        time_dim=1,
    )
    t.assert_close(
        result,
        torch.ones((variable,)),
        msg="RMSE of time mean between zero and 1 should be 1.",
    )
//  [python]
def test_rmse_of_time_mean(variable, time, lats, n_lon):
    x = torch.randn(variable, time, len(lats), n_lon)
    random_weights = torch.rand(len(lats), n_lon)

    result = fme.rmse_of_time_mean(x, x.clone(), time_dim=1)
    torch.testing.assert_close(
        result,
        torch.zeros((variable,)),
        msg="RMSE of time mean between identical tensors should be zero.",
    )

    result = fme.rmse_of_time_mean(
        torch.zeros(variable, time, len(lats), n_lon),
        torch.ones(variable, time, len(lats), n_lon),
        weights=random_weights,
        time_dim=1,
    )
    torch.testing.assert_close(
        result,
        torch.ones((variable,)),
        msg="RMSE of time mean between zero and 1 should be 1.",
    )
//  [python]
def test_weighted_mean_gradient_magnitude():
    constant = torch.ones((5, 2, 4, 4))
    constant_grad_magnitude = fme.weighted_mean_gradient_magnitude(
        constant, dim=(-2, -1)
    )
    t.assert_allclose(constant_grad_magnitude, torch.zeros((5, 2)))

    monotonic = torch.tile(torch.arange(4.0), (5, 2, 4, 1))
    monotonic_grad_magnitude = fme.weighted_mean_gradient_magnitude(
        monotonic, dim=(-2, -1)
    )
    t.assert_allclose(monotonic_grad_magnitude, torch.ones((5, 2)))
//  [python]
def test_weighted_mean_gradient_magnitude():
    constant = torch.ones((5, 2, 4, 4))
    constant_grad_magnitude = fme.weighted_mean_gradient_magnitude(
        constant, dim=(-2, -1)
    )
    torch.testing.assert_close(constant_grad_magnitude, torch.zeros((5, 2)))

    monotonic = torch.tile(torch.arange(4.0), (5, 2, 4, 1))
    monotonic_grad_magnitude = fme.weighted_mean_gradient_magnitude(
        monotonic, dim=(-2, -1)
    )
    torch.testing.assert_close(monotonic_grad_magnitude, torch.ones((5, 2)))
//  [python]
def test_gradient_magnitude_percent_diff():
    constant = torch.ones((5, 2, 4, 4))
    self_percent_diff = fme.gradient_magnitude_percent_diff(
        constant, constant, dim=(-2, -1)
    )
    assert torch.all(torch.isnan(self_percent_diff))

    monotonic = torch.tile(torch.arange(4.0), (5, 2, 4, 1))
    percent_diff = fme.gradient_magnitude_percent_diff(
        monotonic, constant, dim=(-2, -1)
    )
    torch.testing.assert_allclose(percent_diff, -100 * torch.ones((5, 2)))
//  [python]
def test_gradient_magnitude_percent_diff():
    constant = torch.ones((5, 2, 4, 4))
    self_percent_diff = fme.gradient_magnitude_percent_diff(
        constant, constant, dim=(-2, -1)
    )
    assert torch.all(torch.isnan(self_percent_diff))

    monotonic = torch.tile(torch.arange(4.0), (5, 2, 4, 1))
    percent_diff = fme.gradient_magnitude_percent_diff(
        monotonic, constant, dim=(-2, -1)
    )
    torch.testing.assert_close(percent_diff, -100 * torch.ones((5, 2)))
//  [python]
def test_time_and_global_mean(variable, time, lats, n_lon):
    x = torch.randn(variable, time, len(lats), n_lon)
    random_weights = torch.rand(len(lats), n_lon)

    result = fme.time_and_global_mean_bias(x, x.clone(), time_dim=1)
    t.assert_close(
        result,
        torch.zeros((variable,)),
        msg="Time and global mean bias between identical tensors should be zero.",
    )

    result = fme.time_and_global_mean_bias(
        torch.zeros(variable, time, len(lats), n_lon),
        torch.ones(variable, time, len(lats), n_lon),
        weights=random_weights,
        time_dim=1,
    )
    t.assert_close(
        result,
        torch.ones((variable,)),
        msg="Global mean bias between zero and 1 should be 1.",
    )
//  [python]
def test_time_and_global_mean(variable, time, lats, n_lon):
    x = torch.randn(variable, time, len(lats), n_lon)
    random_weights = torch.rand(len(lats), n_lon)

    result = fme.time_and_global_mean_bias(x, x.clone(), time_dim=1)
    torch.testing.assert_close(
        result,
        torch.zeros((variable,)),
        msg="Time and global mean bias between identical tensors should be zero.",
    )

    result = fme.time_and_global_mean_bias(
        torch.zeros(variable, time, len(lats), n_lon),
        torch.ones(variable, time, len(lats), n_lon),
        weights=random_weights,
        time_dim=1,
    )
    torch.testing.assert_close(
        result,
        torch.ones((variable,)),
        msg="Global mean bias between zero and 1 should be 1.",
    )
//  [python]
def test_repeat():
    data = torch.rand((5, 100, 2, 10, 10))
    conv = nn.Conv2d(2, 4, 10)
    repeated_conv = sl.Repeat(conv)
    output = repeated_conv(data)
    alt_output = conv(data.flatten(0, 1)).unflatten(0, (5, 100))

    if version.parse(torch.__version__) > version.parse("1.12"):
        torch.testing.assert_allclose(output, alt_output)
    else:
        torch.testing.assert_allclose(output, alt_output)
//  [python]
def test_repeat():
    data = torch.rand((5, 100, 2, 10, 10))
    conv = nn.Conv2d(2, 4, 10)
    repeated_conv = sl.Repeat(conv)
    output = repeated_conv(data)
    alt_output = conv(data.flatten(0, 1)).unflatten(0, (5, 100))

    if version.parse(torch.__version__) > version.parse("1.12"):
        torch.testing.assert_close(output, alt_output)
    else:
        torch.testing.assert_allclose(output, alt_output)
//  [python]
def test_read_flo(tmpdir):
    path = tmpdir / "test.flo"
    make_fake_flo_file(3, 4, path)

    with open(path, "rb") as file:
        actual = read_flo(file)

    expected = torch.from_numpy(read_flo_ref(path).astype("f4", copy=False))

    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_read_flo(tmpdir):
    path = tmpdir / "test.flo"
    make_fake_flo_file(3, 4, path)

    with open(path, "rb") as file:
        actual = read_flo(file)

    expected = torch.from_numpy(read_flo_ref(path).astype("f4", copy=False))

    torch.testing.assert_close(actual, expected)
//  [python]
def test_fromfile(tmpdir, np_dtype, torch_dtype, byte_order, count, mode):
    path = tmpdir / "data.bin"
    rng = np.random.RandomState(0)
    rng.randn(5 if count == -1 else count + 1).astype(np_dtype).tofile(path)

    for count_ in (-1, count // 2):
        expected = torch.from_numpy(np.fromfile(path, dtype=np_dtype, count=count_).astype(np_dtype[1:]))

        with open(path, mode) as file:
            actual = fromfile(file, dtype=torch_dtype, byte_order=byte_order, count=count_)

        t.assert_close(actual, expected)
//  [python]
def test_fromfile(tmpdir, np_dtype, torch_dtype, byte_order, count, mode):
    path = tmpdir / "data.bin"
    rng = np.random.RandomState(0)
    rng.randn(5 if count == -1 else count + 1).astype(np_dtype).tofile(path)

    for count_ in (-1, count // 2):
        expected = torch.from_numpy(np.fromfile(path, dtype=np_dtype, count=count_).astype(np_dtype[1:]))

        with open(path, mode) as file:
            actual = fromfile(file, dtype=torch_dtype, byte_order=byte_order, count=count_)

        torch.testing.assert_close(actual, expected)
//  [python]
def test_camera_set_world_poses(self):
        """Test camera function to set specific world pose."""
        camera = RayCasterCamera(self.camera_cfg)
        # play sim
        self.sim.reset()

        # convert to torch tensors
        position = torch.tensor([POSITION], dtype=torch.float32, device=camera.device)
        orientation = torch.tensor([QUAT_WORLD], dtype=torch.float32, device=camera.device)
        # set new pose
        camera.set_world_poses(position.clone(), orientation.clone(), convention="world")

        # check if transform correctly set in output
        torch.testing.assert_allclose(camera.data.pos_w, position)
        torch.testing.assert_allclose(camera.data.quat_w_world, orientation)
//  [python]
def test_camera_set_world_poses(self):
        """Test camera function to set specific world pose."""
        camera = RayCasterCamera(self.camera_cfg)
        # play sim
        self.sim.reset()

        # convert to torch tensors
        position = torch.tensor([POSITION], dtype=torch.float32, device=camera.device)
        orientation = torch.tensor([QUAT_WORLD], dtype=torch.float32, device=camera.device)
        # set new pose
        camera.set_world_poses(position.clone(), orientation.clone(), convention="world")

        # check if transform correctly set in output
        torch.testing.assert_close(camera.data.pos_w, position)
        torch.testing.assert_close(camera.data.quat_w_world, orientation)
//  [python]
def test_loss(self, tmpdir, monkeypatch):
        ckpt = "small_test_ckpt_hf"
        ckpt_path = Path(CKPT_MODEL_PATHS[ckpt])
        ckpt_dir = ckpt_path.parent
        log_file = gen_log_file_name(tmpdir)

        # Config file needed for model conversion.
        write_hf_ckpt_config(ckpt_dir)

        cmd = f"""
        tune run --nnodes 1 --nproc_per_node 2 full_finetune_distributed \
            --config llama2/7B_full \
            output_dir={tmpdir} \
            checkpointer._component_=torchtune.utils.FullModelHFCheckpointer \
            checkpointer.checkpoint_dir='{ckpt_dir}' \
            checkpointer.checkpoint_files=[{ckpt_path}]\
            checkpointer.output_dir={tmpdir} \
            checkpointer.model_type=LLAMA2 \
            metric_logger.filename={log_file} \
        """.split()
        model_config = llama2_test_config()
        cmd = cmd + self._get_test_config_overrides() + model_config

        monkeypatch.setattr(sys, "argv", cmd)
        runpy.run_path(TUNE_PATH, run_name="__main__")
        loss_values = get_loss_values_from_metric_logger(log_file)
        expected_loss_values = self._fetch_expected_loss_values(ckpt)
        t.assert_allclose(
            loss_values, expected_loss_values, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_loss(self, tmpdir, monkeypatch):
        ckpt = "small_test_ckpt_hf"
        ckpt_path = Path(CKPT_MODEL_PATHS[ckpt])
        ckpt_dir = ckpt_path.parent
        log_file = gen_log_file_name(tmpdir)

        # Config file needed for model conversion.
        write_hf_ckpt_config(ckpt_dir)

        cmd = f"""
        tune run --nnodes 1 --nproc_per_node 2 full_finetune_distributed \
            --config llama2/7B_full \
            output_dir={tmpdir} \
            checkpointer._component_=torchtune.utils.FullModelHFCheckpointer \
            checkpointer.checkpoint_dir='{ckpt_dir}' \
            checkpointer.checkpoint_files=[{ckpt_path}]\
            checkpointer.output_dir={tmpdir} \
            checkpointer.model_type=LLAMA2 \
            metric_logger.filename={log_file} \
        """.split()
        model_config = llama2_test_config()
        cmd = cmd + self._get_test_config_overrides() + model_config

        monkeypatch.setattr(sys, "argv", cmd)
        runpy.run_path(TUNE_PATH, run_name="__main__")
        loss_values = get_loss_values_from_metric_logger(log_file)
        expected_loss_values = self._fetch_expected_loss_values(ckpt)
        torch.testing.assert_close(
            loss_values, expected_loss_values, rtol=1e-4, atol=1e-4
        )
//  [python]
def test_camera_set_world_poses_from_view(self):
        """Test camera function to set specific world pose from view."""
        camera = RayCasterCamera(self.camera_cfg)
        # play sim
        self.sim.reset()

        # convert to torch tensors
        eyes = torch.tensor([POSITION], dtype=torch.float32, device=camera.device)
        targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera.device)
        quat_ros_gt = torch.tensor([QUAT_ROS], dtype=torch.float32, device=camera.device)
        # set new pose
        camera.set_world_poses_from_view(eyes.clone(), targets.clone())

        # check if transform correctly set in output
        torch.allclose(camera.data.pos_w, eyes)
        torch.allclose(camera.data.quat_w_ros, quat_ros_gt)
//  [python]
def test_camera_set_world_poses_from_view(self):
        """Test camera function to set specific world pose from view."""
        camera = RayCasterCamera(self.camera_cfg)
        # play sim
        self.sim.reset()

        # convert to torch tensors
        eyes = torch.tensor([POSITION], dtype=torch.float32, device=camera.device)
        targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera.device)
        quat_ros_gt = torch.tensor([QUAT_ROS], dtype=torch.float32, device=camera.device)
        # set new pose
        camera.set_world_poses_from_view(eyes.clone(), targets.clone())

        # check if transform correctly set in output
        torch.testing.assert_close(camera.data.pos_w, eyes)
        torch.testing.assert_close(camera.data.quat_w_ros, quat_ros_gt)
//  [python]
def test_rms_norm(self):
        bsz = 6
        s = 18
        hidden_size = 4096
        # -----------golden-----------

        input = torch.rand(bsz, s, hidden_size).cuda()
        input.requires_grad_()
        input.retain_grad()
        non_parallel_norm = get_model().cuda()
        golden_output = non_parallel_norm(input)
        golden_loss = golden_output.mean()
        golden_loss.backward()

        # -----------vescale----------
        device_mesh = DeviceMesh(self.device_type, range(self.world_size))
        vescale_norm = get_model().cuda()
        fwd_resharding_plan = {
            ".input": [[Shard(1)]],
        }
        param_sharding_plan = {}

        vescale_norm = parallelize_module(
            vescale_norm,
            device_mesh,
            {"parameter": param_sharding_plan, "forward": fwd_resharding_plan},
        )
        d_input = distribute_tensor(input.detach(), device_mesh, [Shard(1)])
        d_input.requires_grad_()
        d_input.retain_grad()

        vescale_output = vescale_norm(d_input)
        vescale_output = vescale_output.redistribute(placements=[Replicate()] * device_mesh.ndim)
        vescale_loss = vescale_output.mean()

        vescale_loss.backward()
        vescale_norm.finish_grad_sync()
        d_gard = d_input.grad.redistribute(placements=[Replicate()] * device_mesh.ndim)
        t.assert_close(vescale_output._local_tensor, golden_output)
        t.assert_close(vescale_loss._local_tensor, golden_loss)
        t.assert_close(d_gard._local_tensor, input.grad)
        t.assert_close(vescale_norm.weight.grad._local_tensor, non_parallel_norm.weight.grad)
//  [python]
def test_rms_norm(self):
        bsz = 6
        s = 18
        hidden_size = 4096
        # -----------golden-----------

        input = torch.rand(bsz, s, hidden_size).cuda()
        input.requires_grad_()
        input.retain_grad()
        non_parallel_norm = get_model().cuda()
        golden_output = non_parallel_norm(input)
        golden_loss = golden_output.mean()
        golden_loss.backward()

        # -----------vescale----------
        device_mesh = DeviceMesh(self.device_type, range(self.world_size))
        vescale_norm = get_model().cuda()
        fwd_resharding_plan = {
            ".input": [[Shard(1)]],
        }
        param_sharding_plan = {}

        vescale_norm = parallelize_module(
            vescale_norm,
            device_mesh,
            {"parameter": param_sharding_plan, "forward": fwd_resharding_plan},
        )
        d_input = distribute_tensor(input.detach(), device_mesh, [Shard(1)])
        d_input.requires_grad_()
        d_input.retain_grad()

        vescale_output = vescale_norm(d_input)
        vescale_output = vescale_output.redistribute(placements=[Replicate()] * device_mesh.ndim)
        vescale_loss = vescale_output.mean()

        vescale_loss.backward()
        vescale_norm.finish_grad_sync()
        d_gard = d_input.grad.redistribute(placements=[Replicate()] * device_mesh.ndim)
        torch.testing.assert_close(vescale_output._local_tensor, golden_output)
        torch.testing.assert_close(vescale_loss._local_tensor, golden_loss)
        torch.testing.assert_close(d_gard._local_tensor, input.grad)
        torch.testing.assert_close(vescale_norm.weight.grad._local_tensor, non_parallel_norm.weight.grad)
//  [python]
def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
        output = fn()
        baseline_output = baseline_fn()
        accuracy = True
        try:
            if self.mode == Mode.FWD:
                t.assert_close(output, baseline_output)
            elif self.mode == Mode.BWD:
                t.assert_close(output.grad, baseline_output.grad)
            else:
                fwd_output, loss = output
                baseline_fwd_output, baseline_loss = baseline_output
                t.assert_close(fwd_output, baseline_fwd_output)
                t.assert_close(loss.grad, baseline_loss.grad)
        except Exception:
            # either the output tensor or the loss grad tensor does not match
            accuracy = False
        finally:
            return accuracy
//  [python]
def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
        output = fn()
        baseline_output = baseline_fn()
        accuracy = True
        try:
            if self.mode == Mode.FWD:
                torch.testing.assert_close(output, baseline_output)
            elif self.mode == Mode.BWD:
                torch.testing.assert_close(output.grad, baseline_output.grad)
            else:
                fwd_output, loss = output
                baseline_fwd_output, baseline_loss = baseline_output
                torch.testing.assert_close(fwd_output, baseline_fwd_output)
                torch.testing.assert_close(loss.grad, baseline_loss.grad)
        except Exception:
            # either the output tensor or the loss grad tensor does not match
            accuracy = False
        finally:
            return accuracy
//  [python]
def test_intrinsic_matrix(self):
        """Checks that the camera's set and retrieve methods work for intrinsic matrix."""
        camera_cfg = copy.deepcopy(self.camera_cfg)
        camera_cfg.pattern_cfg.height = 240
        camera_cfg.pattern_cfg.width = 320
        camera = RayCasterCamera(camera_cfg)
        # play sim
        self.sim.reset()
        # Desired properties (obtained from realsense camera at 320x240 resolution)
        rs_intrinsic_matrix = [229.31640625, 0.0, 164.810546875, 0.0, 229.826171875, 122.1650390625, 0.0, 0.0, 1.0]
        rs_intrinsic_matrix = torch.tensor(rs_intrinsic_matrix, device=camera.device).reshape(3, 3).unsqueeze(0)
        # Set matrix into simulator
        camera.set_intrinsic_matrices(rs_intrinsic_matrix.clone())
        # Simulate for a few steps
        # note: This is a workaround to ensure that the textures are loaded.
        #   Check "Known Issues" section in the documentation for more details.
        for _ in range(5):
            self.sim.step()
        # Simulate physics
        for _ in range(10):
            # perform rendering
            self.sim.step()
            # update camera
            camera.update(self.dt)
            # Check that matrix is correct
            # TODO: This is not correctly setting all values in the matrix since the
            #       vertical aperture and aperture offsets are not being set correctly
            #       This is a bug in the simulator.
            torch.testing.assert_allclose(rs_intrinsic_matrix[0, 0, 0], camera.data.intrinsic_matrices[0, 0, 0])
//  [python]
def test_intrinsic_matrix(self):
        """Checks that the camera's set and retrieve methods work for intrinsic matrix."""
        camera_cfg = copy.deepcopy(self.camera_cfg)
        camera_cfg.pattern_cfg.height = 240
        camera_cfg.pattern_cfg.width = 320
        camera = RayCasterCamera(camera_cfg)
        # play sim
        self.sim.reset()
        # Desired properties (obtained from realsense camera at 320x240 resolution)
        rs_intrinsic_matrix = [229.31640625, 0.0, 164.810546875, 0.0, 229.826171875, 122.1650390625, 0.0, 0.0, 1.0]
        rs_intrinsic_matrix = torch.tensor(rs_intrinsic_matrix, device=camera.device).reshape(3, 3).unsqueeze(0)
        # Set matrix into simulator
        camera.set_intrinsic_matrices(rs_intrinsic_matrix.clone())
        # Simulate for a few steps
        # note: This is a workaround to ensure that the textures are loaded.
        #   Check "Known Issues" section in the documentation for more details.
        for _ in range(5):
            self.sim.step()
        # Simulate physics
        for _ in range(10):
            # perform rendering
            self.sim.step()
            # update camera
            camera.update(self.dt)
            # Check that matrix is correct
            # TODO: This is not correctly setting all values in the matrix since the
            #       vertical aperture and aperture offsets are not being set correctly
            #       This is a bug in the simulator.
            torch.testing.assert_close(rs_intrinsic_matrix[0, 0, 0], camera.data.intrinsic_matrices[0, 0, 0])
//  [python]
def _test_conversion(
        self,
        input_data: Tensor,
        reference_data: Tensor,
        ebits: int,
        exponent_bias: int,
        max_pos: float,
        atol: float = 0.0,
        rtol: float = 1e-7,
    ) -> None:
        if torch.cuda.is_available():
            input_data_gpu = input_data.cuda()
            quantized_data_gpu = torch.ops.fbgemm.FloatToHFP8Quantized(
                input_data_gpu, ebits, exponent_bias, max_pos
            )
            dequantized_data_gpu = torch.ops.fbgemm.HFP8QuantizedToFloat(
                quantized_data_gpu, ebits, exponent_bias
            )
            torch.testing.assert_allclose(
                dequantized_data_gpu.cpu(), reference_data, rtol=rtol, atol=atol
            )
//  [python]
def _test_conversion(
        self,
        input_data: Tensor,
        reference_data: Tensor,
        ebits: int,
        exponent_bias: int,
        max_pos: float,
        atol: float = 0.0,
        rtol: float = 1e-7,
    ) -> None:
        if torch.cuda.is_available():
            input_data_gpu = input_data.cuda()
            quantized_data_gpu = torch.ops.fbgemm.FloatToHFP8Quantized(
                input_data_gpu, ebits, exponent_bias, max_pos
            )
            dequantized_data_gpu = torch.ops.fbgemm.HFP8QuantizedToFloat(
                quantized_data_gpu, ebits, exponent_bias
            )
            torch.testing.assert_close(
                dequantized_data_gpu.cpu(), reference_data, rtol=rtol, atol=atol
            )
//  [python]
def test_log_cdf():
    rng = numpy.random.default_rng(0xFEEEFEEE)

    x = rng.random((32,), dtype=numpy.float64)

    actual = log_cdf(torch.tensor(x, dtype=torch.float), Normal(0.0, 1.0))

    assert torch.isfinite(actual).all()

    torch.testing.assert_allclose(
        actual,
        torch.tensor(scipy.special.log_ndtr(x), dtype=torch.float),
    )

    actual = log_cdf(torch.tensor(x, dtype=torch.float), Uniform(0.0, 1.0))

    assert torch.isfinite(actual).all()

    torch.testing.assert_allclose(
        actual,
        torch.tensor(scipy.stats.uniform.logcdf(x), dtype=torch.float),
    )
//  [python]
def test_log_cdf():
    rng = numpy.random.default_rng(0xFEEEFEEE)

    x = rng.random((32,), dtype=numpy.float64)

    actual = log_cdf(torch.tensor(x, dtype=torch.float), Normal(0.0, 1.0))

    assert torch.isfinite(actual).all()

    torch.testing.assert_close(
        actual,
        torch.tensor(scipy.special.log_ndtr(x), dtype=torch.float),
    )

    actual = log_cdf(torch.tensor(x, dtype=torch.float), Uniform(0.0, 1.0))

    assert torch.isfinite(actual).all()

    torch.testing.assert_close(
        actual,
        torch.tensor(scipy.stats.uniform.logcdf(x), dtype=torch.float),
    )
//  [python]
def _test_power(self, input_size, exp, test_name, ait_dtype):
        print(f"Running test {test_name} with exp = {exp}")
        assert len(input_size) == 2
        torch_dtype = _AIT_DTYPE_TO_PYTORCH_DTYPE[ait_dtype]
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=ait_dtype,
            name="input0",
            is_input=True,
        )
        X2 = ops.elementwise(FuncEnum.POW)(X1, exp)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", test_name)

        if abs(exp) < 1.0:
            x1_pt = torch.randn(input_size).cuda().to(dtype=torch_dtype) + 0.5
        else:
            x1_pt = torch.randn(input_size).cuda().to(dtype=torch_dtype)
        x2_pt = torch.pow(x1_pt, exp)

        x2 = torch.empty(input_size).cuda().to(dtype=torch_dtype)
        module.run_with_tensors([x1_pt], [x2])
        # t, _, _ = module.benchmark_with_tensors([x1_pt], [x2], count=1000)
        # bw = input_size[0] * input_size[1] * 2 * 2 / (t * 1e9 * 1e-3)
        # print(f"BW: {bw} GB/s")
        torch.testing.assert_allclose(x2, x2_pt, atol=1e-3, rtol=1e-3, equal_nan=True)
//  [python]
def _test_power(self, input_size, exp, test_name, ait_dtype):
        print(f"Running test {test_name} with exp = {exp}")
        assert len(input_size) == 2
        torch_dtype = _AIT_DTYPE_TO_PYTORCH_DTYPE[ait_dtype]
        X1 = Tensor(
            shape=[IntImm(input_size[0]), IntImm(input_size[1])],
            dtype=ait_dtype,
            name="input0",
            is_input=True,
        )
        X2 = ops.elementwise(FuncEnum.POW)(X1, exp)
        X2._attrs["is_output"] = True
        X2._attrs["name"] = "output0"

        target = detect_target()
        module = compile_model(X2, target, "./tmp", test_name)

        if abs(exp) < 1.0:
            x1_pt = torch.randn(input_size).cuda().to(dtype=torch_dtype) + 0.5
        else:
            x1_pt = torch.randn(input_size).cuda().to(dtype=torch_dtype)
        x2_pt = torch.pow(x1_pt, exp)

        x2 = torch.empty(input_size).cuda().to(dtype=torch_dtype)
        module.run_with_tensors([x1_pt], [x2])
        # t, _, _ = module.benchmark_with_tensors([x1_pt], [x2], count=1000)
        # bw = input_size[0] * input_size[1] * 2 * 2 / (t * 1e9 * 1e-3)
        # print(f"BW: {bw} GB/s")
        torch.testing.assert_close(x2, x2_pt, atol=1e-3, rtol=1e-3, equal_nan=True)
//  [python]
def test_si_sdr_loss(self):
        loss_fn = SISDRLoss()
        sdr_fn = ScaleInvariantSignalDistortionRatio(zero_mean=True)

        num_samples = 1000
        torch.manual_seed(100)
        target = torch.rand([1, num_samples])
        predicted = torch.rand([1, num_samples])
        target_len = torch.tensor([num_samples, num_samples])

        torch_si_sdr = sdr_fn(preds=predicted, target=target)
        loss = loss_fn(audio_real=target, audio_gen=predicted, audio_len=target_len)
        si_sdr = -loss

        torch.testing.assert_allclose(actual=si_sdr, expected=torch_si_sdr)
//  [python]
def test_si_sdr_loss(self):
        loss_fn = SISDRLoss()
        sdr_fn = ScaleInvariantSignalDistortionRatio(zero_mean=True)

        num_samples = 1000
        torch.manual_seed(100)
        target = torch.rand([1, num_samples])
        predicted = torch.rand([1, num_samples])
        target_len = torch.tensor([num_samples, num_samples])

        torch_si_sdr = sdr_fn(preds=predicted, target=target)
        loss = loss_fn(audio_real=target, audio_gen=predicted, audio_len=target_len)
        si_sdr = -loss

        torch.testing.assert_close(actual=si_sdr, expected=torch_si_sdr)
//  [python]
def _test_grouped_auc_helper(
        self,
        labels: torch.Tensor,
        predictions: torch.Tensor,
        weights: torch.Tensor,
        expected_auc: torch.Tensor,
        grouping_keys: Optional[torch.Tensor] = None,
    ) -> None:
        num_task = labels.shape[0]
        batch_size = labels.shape[0]
        task_list = []
        inputs: Dict[str, Union[Dict[str, torch.Tensor], torch.Tensor]] = {
            "predictions": {},
            "labels": {},
            "weights": {},
        }
        if grouping_keys is not None:
            inputs["required_inputs"] = {"grouping_keys": grouping_keys}
        for i in range(num_task):
            task_info = RecTaskInfo(
                name=f"Task:{i}",
                label_name="label",
                prediction_name="prediction",
                weight_name="weight",
            )
            task_list.append(task_info)
            # pyre-ignore
            inputs["predictions"][task_info.name] = predictions[i]
            # pyre-ignore
            inputs["labels"][task_info.name] = labels[i]
            # pyre-ignore
            inputs["weights"][task_info.name] = weights[i]

        auc = AUCMetric(
            world_size=1,
            my_rank=0,
            batch_size=batch_size,
            tasks=task_list,
            # pyre-ignore
            grouped_auc=True,
        )
        # pyre-ignore
        auc.update(**inputs)
        actual_auc = auc.compute()

        for task_id, task in enumerate(task_list):
            cur_actual_auc = actual_auc[f"auc-{task.name}|window_grouped_auc"]
            cur_expected_auc = expected_auc[task_id].unsqueeze(dim=0)

            t.assert_close(
                cur_actual_auc,
                cur_expected_auc,
                atol=1e-4,
                rtol=1e-4,
                check_dtype=False,
                msg=f"Actual: {cur_actual_auc}, Expected: {cur_expected_auc}",
            )
//  [python]
def _test_grouped_auc_helper(
        self,
        labels: torch.Tensor,
        predictions: torch.Tensor,
        weights: torch.Tensor,
        expected_auc: torch.Tensor,
        grouping_keys: Optional[torch.Tensor] = None,
    ) -> None:
        num_task = labels.shape[0]
        batch_size = labels.shape[0]
        task_list = []
        inputs: Dict[str, Union[Dict[str, torch.Tensor], torch.Tensor]] = {
            "predictions": {},
            "labels": {},
            "weights": {},
        }
        if grouping_keys is not None:
            inputs["required_inputs"] = {"grouping_keys": grouping_keys}
        for i in range(num_task):
            task_info = RecTaskInfo(
                name=f"Task:{i}",
                label_name="label",
                prediction_name="prediction",
                weight_name="weight",
            )
            task_list.append(task_info)
            # pyre-ignore
            inputs["predictions"][task_info.name] = predictions[i]
            # pyre-ignore
            inputs["labels"][task_info.name] = labels[i]
            # pyre-ignore
            inputs["weights"][task_info.name] = weights[i]

        auc = AUCMetric(
            world_size=1,
            my_rank=0,
            batch_size=batch_size,
            tasks=task_list,
            # pyre-ignore
            grouped_auc=True,
        )
        # pyre-ignore
        auc.update(**inputs)
        actual_auc = auc.compute()

        for task_id, task in enumerate(task_list):
            cur_actual_auc = actual_auc[f"auc-{task.name}|window_grouped_auc"]
            cur_expected_auc = expected_auc[task_id].unsqueeze(dim=0)

            torch.testing.assert_close(
                cur_actual_auc,
                cur_expected_auc,
                atol=1e-4,
                rtol=1e-4,
                check_dtype=False,
                msg=f"Actual: {cur_actual_auc}, Expected: {cur_expected_auc}",
            )
//  [python]
def test_si_sdr_loss_batch(self):
        loss_fn = SISDRLoss()
        si_sdr_fn = ScaleInvariantSignalDistortionRatio(zero_mean=True)

        batch_size = 3
        num_samples = 1000
        torch.manual_seed(100)
        target = torch.rand([batch_size, num_samples])
        predicted = torch.rand([batch_size, num_samples])

        target_len = torch.tensor([500, 250, 900])
        target = mask_sequence_tensor(target, lengths=target_len)
        predicted = mask_sequence_tensor(predicted, lengths=target_len)

        torch_si_sdr = 0.0
        for i in range(batch_size):
            torch_si_sdr += si_sdr_fn(preds=predicted[i, : target_len[i]], target=target[i, : target_len[i]])
        torch_si_sdr /= batch_size

        loss = loss_fn(audio_real=target, audio_gen=predicted, audio_len=target_len)
        si_sdr = -loss

        self.assertEqual(si_sdr, torch_si_sdr)
//  [python]
def test_si_sdr_loss_batch(self):
        loss_fn = SISDRLoss()
        si_sdr_fn = ScaleInvariantSignalDistortionRatio(zero_mean=True)

        batch_size = 3
        num_samples = 1000
        torch.manual_seed(100)
        target = torch.rand([batch_size, num_samples])
        predicted = torch.rand([batch_size, num_samples])

        target_len = torch.tensor([500, 250, 900])
        target = mask_sequence_tensor(target, lengths=target_len)
        predicted = mask_sequence_tensor(predicted, lengths=target_len)

        torch_si_sdr = 0.0
        for i in range(batch_size):
            torch_si_sdr += si_sdr_fn(preds=predicted[i, : target_len[i]], target=target[i, : target_len[i]])
        torch_si_sdr /= batch_size

        loss = loss_fn(audio_real=target, audio_gen=predicted, audio_len=target_len)
        si_sdr = -loss

        torch.testing.assert_close(actual=si_sdr, expected=torch_si_sdr)
//  [python]
def test_scan2d_broadcast(RBLOCK, num_warps):

    @triton.jit(debug=True)
    def fn(in_ptr, out_ptr, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
        rindex = tl.arange(0, RBLOCK)[None, :]
        xindex = tl.arange(0, XBLOCK)[:, None]
        data = tl.load(in_ptr + rindex)
        scan = tl.cumsum(data, 1)
        expected_max = tl.sum(data, 1)
        tl.device_assert(scan <= expected_max)
        tl.store(out_ptr + xindex * RBLOCK + rindex, scan)

    XBLOCK = 4
    input = torch.randint(0, 10, (1, RBLOCK), dtype=torch.int64, device='cuda')
    output = torch.empty((XBLOCK, RBLOCK), dtype=torch.int64, device='cuda')
    fn[(1, )](input, output, XBLOCK, RBLOCK, num_warps=num_warps)
    ref = input.cumsum(1).broadcast_to((XBLOCK, RBLOCK))
    torch.testing.assert_allclose(output, ref)
//  [python]
def test_scan2d_broadcast(RBLOCK, num_warps):

    @triton.jit(debug=True)
    def fn(in_ptr, out_ptr, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
        rindex = tl.arange(0, RBLOCK)[None, :]
        xindex = tl.arange(0, XBLOCK)[:, None]
        data = tl.load(in_ptr + rindex)
        scan = tl.cumsum(data, 1)
        expected_max = tl.sum(data, 1)
        tl.device_assert(scan <= expected_max)
        tl.store(out_ptr + xindex * RBLOCK + rindex, scan)

    XBLOCK = 4
    input = torch.randint(0, 10, (1, RBLOCK), dtype=torch.int64, device='cuda')
    output = torch.empty((XBLOCK, RBLOCK), dtype=torch.int64, device='cuda')
    fn[(1, )](input, output, XBLOCK, RBLOCK, num_warps=num_warps)
    ref = input.cumsum(1).broadcast_to((XBLOCK, RBLOCK))
    torch.testing.assert_close(output, ref)
//  [python]
def test_pad_constant_weight(self, dtype):
        target = detect_target()
        if dtype == "float" and (int(target._arch) < 80 or target.name == "rocm"):
            self.skipTest("gemm with float tensors requires CUDA sm >= 80")

        M, N, K = 16, 32, 3
        w_pt = get_random_torch_tensor((K, N), dtype)
        input_0 = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[K, N], dtype=dtype, name="weight")
        Y = ops.gemm_rrr()(input_0, W)
        Y._attrs["name"] = "y"
        Y._attrs["is_output"] = True

        mod = compile_model(Y, target, "./tmp", f"test_pad_constant_weight_{dtype}")
        mod.set_constant_with_tensor("weight", w_pt)
        mod.fold_constants()

        input_0_pt = get_random_torch_tensor((M, K), dtype)
        y_pt = torch.matmul(input_0_pt, w_pt)

        y = get_torch_empty_tensor((M, N), dtype)
        mod.run_with_tensors({"input_0": input_0_pt}, {"y": y})

        t.assert_allclose(y, y_pt, atol=1e-1, rtol=1e-1)

        # For float16 inputs, the apply_padding graph pass will add padding to
        # both the input and the weight in this case with concatenate().
        # The concatenate for the weight will be folded, so we will be left with
        # 2 constants.
        if dtype == "float16":
            expected_num_constants = 2
            expected_num_nodes = 5
        elif dtype == "float":
            # Gemm ops with float inputs do not have any alignment requirements,
            # so the apply_padding pass will not add any padding constants.
            # The final graph only contains the original "weight" constant tensor.
            expected_num_constants = 1
            expected_num_nodes = 3
        else:
            raise RuntimeError(f"invalid {dtype=}")
        self._verify_graph(
            mod,
            expected_num_constants=expected_num_constants,
            expected_num_nodes=expected_num_nodes,
        )
//  [python]
def test_pad_constant_weight(self, dtype):
        target = detect_target()
        if dtype == "float" and (int(target._arch) < 80 or target.name == "rocm"):
            self.skipTest("gemm with float tensors requires CUDA sm >= 80")

        M, N, K = 16, 32, 3
        w_pt = get_random_torch_tensor((K, N), dtype)
        input_0 = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[K, N], dtype=dtype, name="weight")
        Y = ops.gemm_rrr()(input_0, W)
        Y._attrs["name"] = "y"
        Y._attrs["is_output"] = True

        mod = compile_model(Y, target, "./tmp", f"test_pad_constant_weight_{dtype}")
        mod.set_constant_with_tensor("weight", w_pt)
        mod.fold_constants()

        input_0_pt = get_random_torch_tensor((M, K), dtype)
        y_pt = torch.matmul(input_0_pt, w_pt)

        y = get_torch_empty_tensor((M, N), dtype)
        mod.run_with_tensors({"input_0": input_0_pt}, {"y": y})

        torch.testing.assert_close(y, y_pt, atol=1e-1, rtol=1e-1)

        # For float16 inputs, the apply_padding graph pass will add padding to
        # both the input and the weight in this case with concatenate().
        # The concatenate for the weight will be folded, so we will be left with
        # 2 constants.
        if dtype == "float16":
            expected_num_constants = 2
            expected_num_nodes = 5
        elif dtype == "float":
            # Gemm ops with float inputs do not have any alignment requirements,
            # so the apply_padding pass will not add any padding constants.
            # The final graph only contains the original "weight" constant tensor.
            expected_num_constants = 1
            expected_num_nodes = 3
        else:
            raise RuntimeError(f"invalid {dtype=}")
        self._verify_graph(
            mod,
            expected_num_constants=expected_num_constants,
            expected_num_nodes=expected_num_nodes,
        )