jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_mbgmm(self, input, paged_lora_a, paged_lora_b, out_head_size,
                   start_loc, seq_lens, adapter_ids, scaling, page_table,
                   ranks, page_start, gt):
        max_seq_len = max(seq_lens).item()
        max_rank = page_table.size(-1)

        xa = mbgmm_a(input,
                     paged_lora_a,
                     q_start_loc=start_loc,
                     q_seqlens=seq_lens,
                     adapter_ids=adapter_ids,
                     rank_page_table=page_table,
                     rank_page_start=page_start,
                     ranks=ranks,
                     max_seq_len=max_seq_len,
                     max_rank=max_rank)

        output = mbgmm_b(xa,
                         paged_lora_b[..., :out_head_size],
                         q_start_loc=start_loc,
                         q_seqlens=seq_lens,
                         adapter_ids=adapter_ids,
                         scaling=scaling,
                         rank_page_table=page_table,
                         rank_page_start=page_start,
                         ranks=ranks,
                         max_seq_len=max_seq_len,
                         max_rank=max_rank)

        torch.testing.assert_close(gt, output)
//  [python]
def _test_rrr(
        self, M, N, K, use_fp16_acc=True, dtype="float16", atol=1e-1, rtol=1e-1
    ):
        target = detect_target(use_fp16_acc=use_fp16_acc)
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(M, name="batch_dim"), K],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[K, N], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rrr_small_nk()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rrr_small_nk_{self.test_count}"
        )

        for m in M:
            X_pt = get_random_torch_tensor([m, K], dtype)
            W_pt = get_random_torch_tensor([K, N], dtype)
            Y_pt = torch.matmul(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = torch.empty_like(Y_pt)
            module.run_with_tensors(inputs, [y])
            if X_pt.nelement() == 0 or W_pt.nelement() == 0:
                pass
            else:
                torch.testing.assert_allclose(Y_pt, y, atol=atol, rtol=rtol)
        self.test_count += 1
//  [python]
def _test_rrr(
        self, M, N, K, use_fp16_acc=True, dtype="float16", atol=1e-1, rtol=1e-1
    ):
        target = detect_target(use_fp16_acc=use_fp16_acc)
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(M, name="batch_dim"), K],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[K, N], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rrr_small_nk()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rrr_small_nk_{self.test_count}"
        )

        for m in M:
            X_pt = get_random_torch_tensor([m, K], dtype)
            W_pt = get_random_torch_tensor([K, N], dtype)
            Y_pt = torch.matmul(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = torch.empty_like(Y_pt)
            module.run_with_tensors(inputs, [y])
            if X_pt.nelement() == 0 or W_pt.nelement() == 0:
                pass
            else:
                torch.testing.assert_close(Y_pt, y, atol=atol, rtol=rtol)
        self.test_count += 1
//  [python]
def _test_label_smoothing_function(self, dtype):
        # Set label smoothing configuration
        smoothing, padding_idx = 0.1, 0
        N, T, H = 128, 74, 32320
        iters = 10
        loss_func = label_smoothing.SoftmaxCrossEntropyLoss.apply

        for i in range(iters):
            logits, labels, half_to_float = self.gen_test_inputs(
                N, T, H, smoothing, padding_idx)

            # Run original softmax cross entropy with label smoothing
            logits.grad = None
            losses = label_smoothing_raw(logits, labels, padding_idx, smoothing)
            loss = losses.sum()
            loss.backward()

            ref_loss = loss.clone().detach()
            ref_grad = logits.grad.clone().detach()

            # Run optimized softmax cross entropy with label smoothing
            logits.grad = None
            losses = loss_func(logits, labels, smoothing, padding_idx, half_to_float)
            loss = losses.sum()
            loss.backward()

            val_loss = loss.clone().detach()
            val_grad = logits.grad.clone().detach()

            # Validate
            self.print_max_diff_elem(ref_grad, val_grad)
            t.assert_close(val_loss, ref_loss)
            t.assert_close(val_grad, ref_grad)
//  [python]
def _test_label_smoothing_function(self, dtype):
        # Set label smoothing configuration
        smoothing, padding_idx = 0.1, 0
        N, T, H = 128, 74, 32320
        iters = 10
        loss_func = label_smoothing.SoftmaxCrossEntropyLoss.apply

        for i in range(iters):
            logits, labels, half_to_float = self.gen_test_inputs(
                N, T, H, smoothing, padding_idx)

            # Run original softmax cross entropy with label smoothing
            logits.grad = None
            losses = label_smoothing_raw(logits, labels, padding_idx, smoothing)
            loss = losses.sum()
            loss.backward()

            ref_loss = loss.clone().detach()
            ref_grad = logits.grad.clone().detach()

            # Run optimized softmax cross entropy with label smoothing
            logits.grad = None
            losses = loss_func(logits, labels, smoothing, padding_idx, half_to_float)
            loss = losses.sum()
            loss.backward()

            val_loss = loss.clone().detach()
            val_grad = logits.grad.clone().detach()

            # Validate
            self.print_max_diff_elem(ref_grad, val_grad)
            torch.testing.assert_close(val_loss, ref_loss)
            torch.testing.assert_close(val_grad, ref_grad)
//  [python]
def test_collate_cut_multi_channel_recording_and_custom_recording_diff_num_channels():
    cut = dummy_multi_cut(0, channel=[0, 1, 2, 3], with_data=True)
    cut.target_recording = dummy_multi_channel_recording(
        1, channel_ids=[0, 1], with_data=True
    )
    cut2 = dummy_multi_cut(2, duration=2.0, channel=[0, 1, 2, 3], with_data=True)
    cut2.target_recording = dummy_multi_channel_recording(
        3, duration=2.0, channel_ids=[0, 1], with_data=True
    )
    cuts = CutSet([cut, cut2])

    expected_lens = torch.tensor([16000, 32000], dtype=torch.int32)

    audio, audio_lens = collate_audio(cuts)
    assert audio.shape == (2, 4, 32000)  # batch x channel x time
    t.assert_close(audio_lens, expected_lens)

    target_audio, target_audio_lens = collate_audio(
        cuts, recording_field="target_recording"
    )
    assert target_audio.shape == (2, 2, 32000)  # batch x channel x time
    t.assert_close(audio_lens, expected_lens)
//  [python]
def test_collate_cut_multi_channel_recording_and_custom_recording_diff_num_channels():
    cut = dummy_multi_cut(0, channel=[0, 1, 2, 3], with_data=True)
    cut.target_recording = dummy_multi_channel_recording(
        1, channel_ids=[0, 1], with_data=True
    )
    cut2 = dummy_multi_cut(2, duration=2.0, channel=[0, 1, 2, 3], with_data=True)
    cut2.target_recording = dummy_multi_channel_recording(
        3, duration=2.0, channel_ids=[0, 1], with_data=True
    )
    cuts = CutSet([cut, cut2])

    expected_lens = torch.tensor([16000, 32000], dtype=torch.int32)

    audio, audio_lens = collate_audio(cuts)
    assert audio.shape == (2, 4, 32000)  # batch x channel x time
    torch.testing.assert_close(audio_lens, expected_lens)

    target_audio, target_audio_lens = collate_audio(
        cuts, recording_field="target_recording"
    )
    assert target_audio.shape == (2, 2, 32000)  # batch x channel x time
    torch.testing.assert_close(audio_lens, expected_lens)
//  [python]
def test_camera_set_world_poses(self):
        """Test camera function to set specific world pose."""
        camera = Camera(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")

        # 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()

        # 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 = Camera(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")

        # 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()

        # 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_grid_clone_env_origins(self):
        """Tests that env origins are consistent when computed using the TerrainImporter and IsaacSim GridCloner."""
        # iterate over different number of environments and environment spacing
        for device in ("cuda:0", "cpu"):
            for env_spacing in [1.0, 4.325, 8.0]:
                for num_envs in [1, 4, 125, 379, 1024]:
                    with self.subTest(num_envs=num_envs, env_spacing=env_spacing):
                        with build_simulation_context(device=device, auto_add_lighting=True) as sim:
                            # create terrain importer
                            terrain_importer_cfg = TerrainImporterCfg(
                                num_envs=num_envs,
                                env_spacing=env_spacing,
                                prim_path="/World/ground",
                                terrain_type="plane",  # for flat ground, origins are in grid
                                terrain_generator=None,
                            )
                            terrain_importer = TerrainImporter(terrain_importer_cfg)
                            # obtain env origins using terrain importer
                            terrain_importer_origins = terrain_importer.env_origins

                            # obtain env origins using grid cloner
                            grid_cloner_origins = self._obtain_grid_cloner_env_origins(
                                num_envs, env_spacing, device=sim.device
                            )

                            # check if the env origins are the same
                            t.assert_close(
                                terrain_importer_origins, grid_cloner_origins, rtol=1e-5, atol=1e-5
                            )
//  [python]
def test_grid_clone_env_origins(self):
        """Tests that env origins are consistent when computed using the TerrainImporter and IsaacSim GridCloner."""
        # iterate over different number of environments and environment spacing
        for device in ("cuda:0", "cpu"):
            for env_spacing in [1.0, 4.325, 8.0]:
                for num_envs in [1, 4, 125, 379, 1024]:
                    with self.subTest(num_envs=num_envs, env_spacing=env_spacing):
                        with build_simulation_context(device=device, auto_add_lighting=True) as sim:
                            # create terrain importer
                            terrain_importer_cfg = TerrainImporterCfg(
                                num_envs=num_envs,
                                env_spacing=env_spacing,
                                prim_path="/World/ground",
                                terrain_type="plane",  # for flat ground, origins are in grid
                                terrain_generator=None,
                            )
                            terrain_importer = TerrainImporter(terrain_importer_cfg)
                            # obtain env origins using terrain importer
                            terrain_importer_origins = terrain_importer.env_origins

                            # obtain env origins using grid cloner
                            grid_cloner_origins = self._obtain_grid_cloner_env_origins(
                                num_envs, env_spacing, device=sim.device
                            )

                            # check if the env origins are the same
                            torch.testing.assert_close(
                                terrain_importer_origins, grid_cloner_origins, rtol=1e-5, atol=1e-5
                            )
//  [python]
def _test_opt_with_module(
        self,
        module: torch.nn.Module,
        inputs: List,
        opt: Callable,
        should_change_graph: bool,
        deleted_ops: Dict = None,
        created_ops: Dict = None,
        rtol: float = None,
        atol: float = None,
    ):
        assert should_change_graph or not bool(deleted_ops or created_ops)
        deleted_ops = deleted_ops or {}
        created_ops = created_ops or {}
        module.eval()

        # Before Opt
        before_results = module(*inputs)
        mod_traced = acc_tracer.trace(module, inputs)
        before_node_list = list(mod_traced.graph.nodes)
        _LOGGER.info("Model before opt.")
        debug_print_graph_module(mod_traced)

        # Apply Opt
        graph_changed = bool(opt(mod_traced))

        # After Opt
        after_results = mod_traced(*inputs)
        after_node_list = list(mod_traced.graph.nodes)
        _LOGGER.info("Model after opt.")
        mod_traced.recompile()
        debug_print_graph_module(mod_traced)

        # Tests
        #  * Numerics
        tol_args = {}
        if rtol is not None:
            tol_args["rtol"] = rtol
        if atol is not None:
            tol_args["atol"] = atol
        t.assert_close(before_results, after_results, **tol_args)

        #  * opt changes graph
        self.assertEqual(graph_changed, before_node_list != after_node_list)
        self.assertEqual(should_change_graph, graph_changed)

        # * modified nodes
        before_node_set = set(before_node_list)
        after_node_set = set(after_node_list)
        self.assertEqual(
            dict(Counter([node.target for node in before_node_set - after_node_set])),
            deleted_ops,
        )
        self.assertEqual(
            dict(Counter([node.target for node in after_node_set - before_node_set])),
            created_ops,
        )

        return mod_traced
//  [python]
def _test_opt_with_module(
        self,
        module: torch.nn.Module,
        inputs: List,
        opt: Callable,
        should_change_graph: bool,
        deleted_ops: Dict = None,
        created_ops: Dict = None,
        rtol: float = None,
        atol: float = None,
    ):
        assert should_change_graph or not bool(deleted_ops or created_ops)
        deleted_ops = deleted_ops or {}
        created_ops = created_ops or {}
        module.eval()

        # Before Opt
        before_results = module(*inputs)
        mod_traced = acc_tracer.trace(module, inputs)
        before_node_list = list(mod_traced.graph.nodes)
        _LOGGER.info("Model before opt.")
        debug_print_graph_module(mod_traced)

        # Apply Opt
        graph_changed = bool(opt(mod_traced))

        # After Opt
        after_results = mod_traced(*inputs)
        after_node_list = list(mod_traced.graph.nodes)
        _LOGGER.info("Model after opt.")
        mod_traced.recompile()
        debug_print_graph_module(mod_traced)

        # Tests
        #  * Numerics
        tol_args = {}
        if rtol is not None:
            tol_args["rtol"] = rtol
        if atol is not None:
            tol_args["atol"] = atol
        torch.testing.assert_close(before_results, after_results, **tol_args)

        #  * opt changes graph
        self.assertEqual(graph_changed, before_node_list != after_node_list)
        self.assertEqual(should_change_graph, graph_changed)

        # * modified nodes
        before_node_set = set(before_node_list)
        after_node_set = set(after_node_list)
        self.assertEqual(
            dict(Counter([node.target for node in before_node_set - after_node_set])),
            deleted_ops,
        )
        self.assertEqual(
            dict(Counter([node.target for node in after_node_set - before_node_set])),
            created_ops,
        )

        return mod_traced
//  [python]
def test_camera_set_world_poses_from_view(self):
        """Test camera function to set specific world pose from view."""
        camera = Camera(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())

        # 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()

        # 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 = Camera(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())

        # 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()

        # 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_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.height = 240
        camera_cfg.width = 320
        camera = Camera(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.height = 240
        camera_cfg.width = 320
        camera = Camera(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_neighbor_list(self, neighbor_list, environment):
        cutoff, props, neighbors_ref = environment
        neighbor_list = neighbor_list(cutoff)
        neighbors = neighbor_list(props)
        R = props[structure.R]
        neighbors[structure.Rij] = (
            R[neighbors[structure.idx_j]]
            - R[neighbors[structure.idx_i]]
            + props[structure.offsets]
        )

        neighbors = self._sort_neighbors(neighbors)
        neighbors_ref = self._sort_neighbors(neighbors_ref)

        for nbl, nbl_ref in zip(neighbors, neighbors_ref):
            t.assert_close(nbl, nbl_ref)
//  [python]
def test_neighbor_list(self, neighbor_list, environment):
        cutoff, props, neighbors_ref = environment
        neighbor_list = neighbor_list(cutoff)
        neighbors = neighbor_list(props)
        R = props[structure.R]
        neighbors[structure.Rij] = (
            R[neighbors[structure.idx_j]]
            - R[neighbors[structure.idx_i]]
            + props[structure.offsets]
        )

        neighbors = self._sort_neighbors(neighbors)
        neighbors_ref = self._sort_neighbors(neighbors_ref)

        for nbl, nbl_ref in zip(neighbors, neighbors_ref):
            torch.testing.assert_close(nbl, nbl_ref)
//  [python]
def test_remove_com(four_atoms):
    positions_trans = SubtractCenterOfMass()(four_atoms)

    com = torch.tensor([0.0, 0.0, 0.0])
    for r_i, m_i in zip(
        positions_trans[structure.position], atomic_masses[four_atoms[structure.Z]]
    ):
        com += r_i * m_i

    torch.allclose(com, torch.tensor([0.0, 0.0, 0.0]))
//  [python]
def test_remove_com(four_atoms):
    positions_trans = SubtractCenterOfMass()(four_atoms)

    com = torch.tensor([0.0, 0.0, 0.0])
    for r_i, m_i in zip(
        positions_trans[structure.position], atomic_masses[four_atoms[structure.Z]]
    ):
        com += r_i * m_i

    torch.testing.assert_close(com, torch.tensor([0.0, 0.0, 0.0]))
//  [python]
def test_remove_cog(four_atoms):
    positions_trans = SubtractCenterOfGeometry()(four_atoms)

    cog = torch.tensor([0.0, 0.0, 0.0])
    for r_i in positions_trans[structure.position]:
        cog += r_i

    torch.allclose(cog, torch.tensor([0.0, 0.0, 0.0]))
//  [python]
def test_remove_cog(four_atoms):
    positions_trans = SubtractCenterOfGeometry()(four_atoms)

    cog = torch.tensor([0.0, 0.0, 0.0])
    for r_i in positions_trans[structure.position]:
        cog += r_i

    torch.testing.assert_close(cog, torch.tensor([0.0, 0.0, 0.0]))
//  [python]
shapes = [
    scalar,
    small,
    small_2d,
    small_broadcast,
    medium,
    medium2,
    medium3d,
    medium_sliced,
    medium_transpose,
    medium_channels_last,
    medium_broadcast,
    medium_broadcast_channels_last,
    large,
    large_transpose,
    large_channels_last,
    pathological_broadcast,
]

operators = [
    add,
    sub,
    mul,
    div,
    relu,
    sigmoid,
    tanh,
    log,
    exp,
    square,
    fma,
    hardswish,
    native_hardswish,
]

nope = set()
for shape, operator in itertools.product(shapes, operators):
    nargs = len(inspect.signature(operator).parameters)
    args = shape()[:nargs]

    try:
        if shape == medium_transpose:
            raise RuntimeError("pointwise_operator hangs on medium_transpose")
        pw_op = pointwise_operator(operator)
        t.assert_close(operator(*args), pw_op(*args))
    except Exception:
        print(f"pointwise_operator failed on {operator.__name__}, {shape.__name__}")
        nope.add((operator, shape))

    ts_op = torch.jit.script(operator)
    t.assert_close(operator(*args), ts_op(*args))
//  [python]
shapes = [
    scalar,
    small,
    small_2d,
    small_broadcast,
    medium,
    medium2,
    medium3d,
    medium_sliced,
    medium_transpose,
    medium_channels_last,
    medium_broadcast,
    medium_broadcast_channels_last,
    large,
    large_transpose,
    large_channels_last,
    pathological_broadcast,
]

operators = [
    add,
    sub,
    mul,
    div,
    relu,
    sigmoid,
    tanh,
    log,
    exp,
    square,
    fma,
    hardswish,
    native_hardswish,
]

nope = set()
for shape, operator in itertools.product(shapes, operators):
    nargs = len(inspect.signature(operator).parameters)
    args = shape()[:nargs]

    try:
        if shape == medium_transpose:
            raise RuntimeError("pointwise_operator hangs on medium_transpose")
        pw_op = pointwise_operator(operator)
        torch.testing.assert_close(operator(*args), pw_op(*args))
    except Exception:
        print(f"pointwise_operator failed on {operator.__name__}, {shape.__name__}")
        nope.add((operator, shape))

    ts_op = torch.jit.script(operator)
    torch.testing.assert_close(operator(*args), ts_op(*args))
//  [python]
def test_gradient_clipping_by_value(tmp_path, precision):
    """Test gradient clipping by value."""
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_steps=1,
        max_epochs=1,
        accelerator="auto",
        devices=1,
        precision=precision,
        gradient_clip_algorithm="value",
        gradient_clip_val=1e-10,
    )

    class TestModel(BoringModel):
        def configure_gradient_clipping(self, *args, **kwargs):
            super().configure_gradient_clipping(*args, **kwargs)
            # test that gradient is clipped correctly
            parameters = self.parameters()
            grad_max_list = [torch.max(p.grad.detach().abs()) for p in parameters]
            grad_max = torch.max(torch.stack(grad_max_list))
            torch.testing.assert_allclose(grad_max.abs(), torch.tensor(1e-10, device=self.device))
            self.assertion_called = True

    model = TestModel()
    trainer.fit(model)
    assert model.assertion_called
//  [python]
def test_gradient_clipping_by_value(tmp_path, precision):
    """Test gradient clipping by value."""
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_steps=1,
        max_epochs=1,
        accelerator="auto",
        devices=1,
        precision=precision,
        gradient_clip_algorithm="value",
        gradient_clip_val=1e-10,
    )

    class TestModel(BoringModel):
        def configure_gradient_clipping(self, *args, **kwargs):
            super().configure_gradient_clipping(*args, **kwargs)
            # test that gradient is clipped correctly
            parameters = self.parameters()
            grad_max_list = [torch.max(p.grad.detach().abs()) for p in parameters]
            grad_max = torch.max(torch.stack(grad_max_list))
            torch.testing.assert_close(grad_max.abs(), torch.tensor(1e-10, device=self.device))
            self.assertion_called = True

    model = TestModel()
    trainer.fit(model)
    assert model.assertion_called
//  [python]
def _test_gemm_rcr_bias_swish(
        self,
        M=128,
        K=1024,
        N=64,
        dtype="float16",
        test_suffix=None,
    ):
        X = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=dtype, name="input_1", is_input=True)
        B = Tensor(shape=[N], dtype=dtype, name="input_2", is_input=True)
        OP = ops.gemm_rcr_bias_swish()
        Y = OP(X, W, B)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        if test_suffix is None:
            test_suffix = dtype
        test_name = f"gemm_rcr_bias_swish_{test_suffix}_{self._test_id}"
        self._test_id += 1
        module = compile_model(Y, detect_target(), "./tmp", test_name)
        X_pt = get_random_torch_tensor([M, K], dtype)
        W_pt = get_random_torch_tensor([N, K], dtype)
        B_pt = get_random_torch_tensor([N], dtype)
        Y_pt = torch.nn.functional.linear(X_pt, W_pt, bias=B_pt)
        Y_pt = swish(Y_pt)

        inputs = {"input_0": X_pt, "input_1": W_pt, "input_2": B_pt}
        y = get_torch_empty_tensor([M, N], dtype)
        module.run_with_tensors(inputs, [y])
        torch.testing.assert_allclose(Y_pt, y, **_TOLERANCE_LIMITS[dtype])
//  [python]
def _test_gemm_rcr_bias_swish(
        self,
        M=128,
        K=1024,
        N=64,
        dtype="float16",
        test_suffix=None,
    ):
        X = Tensor(shape=[M, K], dtype=dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=dtype, name="input_1", is_input=True)
        B = Tensor(shape=[N], dtype=dtype, name="input_2", is_input=True)
        OP = ops.gemm_rcr_bias_swish()
        Y = OP(X, W, B)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        if test_suffix is None:
            test_suffix = dtype
        test_name = f"gemm_rcr_bias_swish_{test_suffix}_{self._test_id}"
        self._test_id += 1
        module = compile_model(Y, detect_target(), "./tmp", test_name)
        X_pt = get_random_torch_tensor([M, K], dtype)
        W_pt = get_random_torch_tensor([N, K], dtype)
        B_pt = get_random_torch_tensor([N], dtype)
        Y_pt = torch.nn.functional.linear(X_pt, W_pt, bias=B_pt)
        Y_pt = swish(Y_pt)

        inputs = {"input_0": X_pt, "input_1": W_pt, "input_2": B_pt}
        y = get_torch_empty_tensor([M, N], dtype)
        module.run_with_tensors(inputs, [y])
        torch.testing.assert_close(Y_pt, y, **_TOLERANCE_LIMITS[dtype])
//  [python]
def test_expand_shrink(self, img, factor):
        # expand then shrink will be the same as the original image, up to this
        # fudge factor
        img = po.load_images(IMG_DIR / "256" / f"{img}.pgm").to(DEVICE)
        modified = po.tools.shrink(po.tools.expand(img, factor), factor)
        t.assert_allclose(img, modified, atol=2e-2, rtol=1e-6)
//  [python]
def test_expand_shrink(self, img, factor):
        # expand then shrink will be the same as the original image, up to this
        # fudge factor
        img = po.load_images(IMG_DIR / "256" / f"{img}.pgm").to(DEVICE)
        modified = po.tools.shrink(po.tools.expand(img, factor), factor)
        torch.testing.assert_close(img, modified, atol=2e-2, rtol=1e-6)
//  [python]
def test_inductor_cummax_bool(device):

    @triton.jit
    def triton_(in_ptr0, out_ptr0, out_ptr1, XBLOCK: tl.constexpr):
        offset = tl.arange(0, XBLOCK)
        tmp0 = tl.load(in_ptr0 + offset).to(tl.int1)
        tmp1 = tmp0.to(tl.int1)
        tmp3 = offset.to(tl.int64)
        tmp5, tmp6, = tl.associative_scan((
            tmp1,
            tmp3,
        ), 0, _triton_cummax_helper_fn)
        tl.store(out_ptr0 + offset, tmp5)
        tl.store(out_ptr1 + offset, tmp6)

    a = torch.randn((64, ), device=device) > 0
    values = torch.empty((64, ), dtype=torch.bool, device=device)
    indices = torch.empty((64, ), dtype=torch.int64, device=device)
    ref = torch.cummax(a, dim=0)

    triton_[(1, )](a, values, indices, 64)
    torch.testing.assert_allclose(ref.values, values)
    torch.testing.assert_allclose(ref.indices, indices)
//  [python]
def test_inductor_cummax_bool(device):

    @triton.jit
    def triton_(in_ptr0, out_ptr0, out_ptr1, XBLOCK: tl.constexpr):
        offset = tl.arange(0, XBLOCK)
        tmp0 = tl.load(in_ptr0 + offset).to(tl.int1)
        tmp1 = tmp0.to(tl.int1)
        tmp3 = offset.to(tl.int64)
        tmp5, tmp6, = tl.associative_scan((
            tmp1,
            tmp3,
        ), 0, _triton_cummax_helper_fn)
        tl.store(out_ptr0 + offset, tmp5)
        tl.store(out_ptr1 + offset, tmp6)

    a = torch.randn((64, ), device=device) > 0
    values = torch.empty((64, ), dtype=torch.bool, device=device)
    indices = torch.empty((64, ), dtype=torch.int64, device=device)
    ref = torch.cummax(a, dim=0)

    triton_[(1, )](a, values, indices, 64)
    torch.testing.assert_close(ref.values, values)
    torch.testing.assert_close(ref.indices, indices)
//  [python]
def test_bfp16_quantize(self):
        X = torch.randn(5 , 10)
        quantized_X = X.to(torch.bfloat16)
        dedequantized_X = quantized_X.to(torch.float32)
        torch.testing.assert_allclose(X, dedequantized_X, rtol=1e-4, atol=5e-3)
//  [python]
def test_bfp16_quantize(self):
        X = torch.randn(5 , 10)
        quantized_X = X.to(torch.bfloat16)
        dedequantized_X = quantized_X.to(torch.float32)
        torch.testing.assert_close(X, dedequantized_X, rtol=1e-4, atol=5e-3)
//  [python]
def test_add_mlp_layer(num_inputs, num_outputs, hidden_size, device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=num_inputs,
        num_outputs=num_outputs,
        hidden_size=hidden_size,
        max_hidden_layers=10,
        device=device,
    )

    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.add_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size < 10:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_allclose(param, initial_net_dict[key])
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_add_mlp_layer(num_inputs, num_outputs, hidden_size, device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=num_inputs,
        num_outputs=num_outputs,
        hidden_size=hidden_size,
        max_hidden_layers=10,
        device=device,
    )

    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.add_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size < 10:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_beam_sample_generate(self):
        for model_class in self.all_generative_model_classes:
            config, input_ids, attention_mask = self._get_input_ids_and_config()

            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])

            model = model_class(config).to(torch_device).eval()
            beam_kwargs = self._get_beam_kwargs()

            output_generate = self._beam_sample_generate(
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
            )

            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])

            if "inputs_embeds" in set(inspect.signature(model.prepare_inputs_for_generation).parameters):
                input_embeds = model.get_input_embeddings()(input_ids)
                beam_kwargs.update({"inputs_embeds": input_embeds})
                output_generate2 = self._beam_sample_generate(
                    model=model,
                    input_ids=None,
                    attention_mask=attention_mask,
                    beam_kwargs=beam_kwargs,
                    logits_warper_kwargs=logits_warper_kwargs,
                )

                torch.testing.assert_allclose(output_generate[:, input_embeds.shape[1] :], output_generate2)
//  [python]
def test_beam_sample_generate(self):
        for model_class in self.all_generative_model_classes:
            config, input_ids, attention_mask = self._get_input_ids_and_config()

            _, logits_warper_kwargs = self._get_logits_processor_and_warper_kwargs(input_ids.shape[-1])

            model = model_class(config).to(torch_device).eval()
            beam_kwargs = self._get_beam_kwargs()

            output_generate = self._beam_sample_generate(
                model=model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                beam_kwargs=beam_kwargs,
                logits_warper_kwargs=logits_warper_kwargs,
            )

            if model.config.is_encoder_decoder:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + 1)
            else:
                self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + input_ids.shape[-1])

            if "inputs_embeds" in set(inspect.signature(model.prepare_inputs_for_generation).parameters):
                input_embeds = model.get_input_embeddings()(input_ids)
                beam_kwargs.update({"inputs_embeds": input_embeds})
                output_generate2 = self._beam_sample_generate(
                    model=model,
                    input_ids=None,
                    attention_mask=attention_mask,
                    beam_kwargs=beam_kwargs,
                    logits_warper_kwargs=logits_warper_kwargs,
                )

                torch.testing.assert_close(output_generate[:, input_embeds.shape[1] :], output_generate2)
//  [python]
def _test_tensor_dict_sync_state(dst_rank: Optional[int] = None) -> None:
    device = init_from_env()

    if dist.get_rank() == 0:
        state_data = {
            _METRIC_NAME: {
                "mapping": {
                    "a": torch.tensor(1, device=device),
                    "b": torch.tensor(10, device=device),
                },
            }
        }
    else:
        state_data = {
            _METRIC_NAME: {
                "mapping": {
                    "a": torch.tensor(2, device=device),
                    "b": torch.tensor(20, device=device),
                },
            }
        }

    # pyre-ignore: Incompatible parameter type [6]:
    dict_items = metrics_traversal_order(state_data)
    synced_states = sync_states(
        state_data, {_METRIC_NAME: device}, dict_items, rank=dst_rank
    )

    tc = unittest.TestCase()
    if dst_rank is None or dist.get_rank() == dst_rank:
        tc.assertIsNotNone(synced_states)
        tc.assertEqual(len(synced_states), 2)

        tc.assertTrue(torch.allclose(
            synced_states[0][_METRIC_NAME]["mapping"]["a"],
            torch.tensor(1, device=device),
        ))
        tc.assertTrue(torch.allclose(
            synced_states[1][_METRIC_NAME]["mapping"]["a"],
            torch.tensor(2, device=device),
        ))

        tc.assertTrue(torch.allclose(
            synced_states[0][_METRIC_NAME]["mapping"]["b"],
            torch.tensor(10, device=device),
        ))
        tc.assertTrue(torch.allclose(
            synced_states[1][_METRIC_NAME]["mapping"]["b"],
            torch.tensor(20, device=device),
        ))
    else:
        tc.assertIsNone(synced_states)
//  [python]
def _test_tensor_dict_sync_state(dst_rank: Optional[int] = None) -> None:
    device = init_from_env()

    if dist.get_rank() == 0:
        state_data = {
            _METRIC_NAME: {
                "mapping": {
                    "a": torch.tensor(1, device=device),
                    "b": torch.tensor(10, device=device),
                },
            }
        }
    else:
        state_data = {
            _METRIC_NAME: {
                "mapping": {
                    "a": torch.tensor(2, device=device),
                    "b": torch.tensor(20, device=device),
                },
            }
        }

    # pyre-ignore: Incompatible parameter type [6]:
    dict_items = metrics_traversal_order(state_data)
    synced_states = sync_states(
        state_data, {_METRIC_NAME: device}, dict_items, rank=dst_rank
    )

    tc = unittest.TestCase()
    if dst_rank is None or dist.get_rank() == dst_rank:
        tc.assertIsNotNone(synced_states)
        tc.assertEqual(len(synced_states), 2)

        torch.testing.assert_close(
            synced_states[0][_METRIC_NAME]["mapping"]["a"],
            torch.tensor(1, device=device),
        )
        torch.testing.assert_close(
            synced_states[1][_METRIC_NAME]["mapping"]["a"],
            torch.tensor(2, device=device),
        )

        torch.testing.assert_close(
            synced_states[0][_METRIC_NAME]["mapping"]["b"],
            torch.tensor(10, device=device),
        )
        torch.testing.assert_close(
            synced_states[1][_METRIC_NAME]["mapping"]["b"],
            torch.tensor(20, device=device),
        )
    else:
        tc.assertIsNone(synced_states)
//  [python]
def test_clone_instance(num_inputs, num_outputs, hidden_size, device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=num_inputs,
        num_outputs=num_outputs,
        hidden_size=hidden_size,
        device=device,
    )
    original_net_dict = dict(evolvable_mlp.feature_net.named_parameters())
    clone = evolvable_mlp.clone()
    clone_net = clone.feature_net
    assert isinstance(clone, EvolvableMLP)
    assert clone.init_dict == evolvable_mlp.init_dict
    assert str(clone.state_dict()) == str(evolvable_mlp.state_dict())
    for key, param in clone_net.named_parameters():
        torch.testing.assert_allclose(param, original_net_dict[key]), evolvable_mlp
//  [python]
def test_clone_instance(num_inputs, num_outputs, hidden_size, device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=num_inputs,
        num_outputs=num_outputs,
        hidden_size=hidden_size,
        device=device,
    )
    original_net_dict = dict(evolvable_mlp.feature_net.named_parameters())
    clone = evolvable_mlp.clone()
    clone_net = clone.feature_net
    assert isinstance(clone, EvolvableMLP)
    assert clone.init_dict == evolvable_mlp.init_dict
    assert str(clone.state_dict()) == str(evolvable_mlp.state_dict())
    for key, param in clone_net.named_parameters():
        torch.testing.assert_close(param, original_net_dict[key]), evolvable_mlp
//  [python]
def test_add_mlp_layer_rainbow(device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=10,
        num_outputs=4,
        hidden_size=[32, 32],
        rainbow=True,
        support=torch.linspace(0.0, 200.0, 50).to(device),
        device=device,
    )
    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.add_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size < 10:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                t.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_add_mlp_layer_rainbow(device):
    evolvable_mlp = EvolvableMLP(
        num_inputs=10,
        num_outputs=4,
        hidden_size=[32, 32],
        rainbow=True,
        support=torch.linspace(0.0, 200.0, 50).to(device),
        device=device,
    )
    initial_hidden_size = len(evolvable_mlp.hidden_size)
    initial_net = evolvable_mlp.feature_net
    initial_net_dict = dict(initial_net.named_parameters())
    evolvable_mlp.add_mlp_layer()
    new_net = evolvable_mlp.feature_net
    if initial_hidden_size < 10:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size + 1
        for key, param in new_net.named_parameters():
            if key in initial_net_dict.keys():
                torch.testing.assert_close(param, initial_net_dict[key])
    else:
        assert len(evolvable_mlp.hidden_size) == initial_hidden_size
//  [python]
def test_fx2ait_args(self):
        class TestModule(torch.nn.Module):
            def forward(self, a, b, c, d):
                temp = a + b[0] + b[1] + c[0] + c[1] + d
                return temp

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

        a = torch.randn(5, 3).half().cuda()
        b = [torch.randn(5, 3).half().cuda(), torch.randn(5, 3).half().cuda()]
        c = [torch.randn(5, 3).half().cuda(), torch.randn(5, 3).half().cuda()]
        d = torch.randn(5, 3).half().cuda()
        ref_output = mod(a, b, c, d)

        traced = acc_tracer.trace(mod, [a, b, c, d])

        ait_dump_dir = tempfile.mkdtemp(prefix="test_fx2ait_", dir="/tmp")

        interp = AITInterpreter(traced, [a, b, c, d], ait_dump_dir, "test")
        interp_result = interp.run()
        ait_mod = AITModule(
            AIT_MODEL_CLASS(
                interp_result.engine.lib_path,
                interp_result.input_names,
                interp_result.output_names,
                torch.float16,
                torch.float16,
                1,  # num_runtimes
            ),
            interp_result,
        )
        ait_output = ait_mod(a, b, c, d)
        self.assertEqual(ait_output, ref_output)

        ait_output = ait_mod(a, b, c, d=d)
        self.assertEqual(ait_output, ref_output)

        ait_output = ait_mod(a, b, c=c, d=d)
        self.assertEqual(ait_output, ref_output)

        ait_output = ait_mod(a, b=b, c=c, d=d)
        self.assertEqual(ait_output, ref_output)