jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check_lie_group_function(
    module, op_name: str, atol: float, inputs, funcs=None, batch_size=None
):
    op_impl = getattr(module, f"_{op_name}_impl")
    op = getattr(module, f"_{op_name}_autograd_fn")
    jop = getattr(module, f"_j{op_name}_autograd_fn")

    # Check jacobians
    jacs_impl = torch.autograd.functional.jacobian(op_impl, inputs, vectorize=True)
    jacs = torch.autograd.functional.jacobian(op, inputs, vectorize=True)

    if funcs is None:
        for jac_impl, jac in zip(jacs_impl, jacs):
            torch.testing.assert_close(jac_impl, jac, atol=atol, rtol=atol)
    else:
        for jac_impl, jac, func in zip(jacs_impl, jacs, funcs):
            if func is None:
                torch.testing.assert_close(jac_impl, jac, atol=atol, rtol=atol)
            else:
                torch.testing.assert_close(
                    func(jac_impl), func(jac), atol=atol, rtol=atol
                )

    # Check multi-batch consistency
    if batch_size is None:
        return
    lb = len(batch_size)
    flattened_inputs = [x.reshape(-1, *x.shape[lb:]) for x in inputs]
    out = op(*inputs)
    flattened_out = op(*flattened_inputs)
    if jop is None:
        return
    jout = jop(*inputs)[0]
    flattened_jout = jop(*flattened_inputs)[0]
    torch.testing.assert_close(out, flattened_out.reshape(*batch_size, *out.shape[lb:]))
    for j, jf in zip(jout, flattened_jout):
        torch.testing.assert_close(j, jf.reshape(*batch_size, *j.shape[lb:]))
//  [python]
def check_jacrev_binary(group_fns, batch_size, name):
    assert name in ["compose", "transform", "untransform"]
    if not hasattr(torch, "vmap"):
        return

    test_fn = getattr(group_fns, name)

    fn_inputs = (
        (group_fns.rand(batch_size), torch.randn(batch_size, 3))
        if name in ["transform", "untransform"]
        else (group_fns.rand(batch_size), group_fns.rand(batch_size))
    )

    def f(t1, t2):
        op_out = test_fn(t1.unsqueeze(0), t2.unsqueeze(0))
        if name == "compose":
            op_out = group_fns.log(op_out)
        return op_out.squeeze(0)

    jacs_vmap = []
    for i in range(2):
        j = torch.vmap(torch.func.jacrev(f, i))(fn_inputs[0], fn_inputs[1])
        if fn_inputs[i].ndim == 3:  # group input
            j = group_fns.left_project(fn_inputs[i], j)
        jacs_vmap.append(j)

    jlog, jtest = [], []
    out = test_fn(fn_inputs[0], fn_inputs[1], jacobians=jtest)
    if name == "compose":
        group_fns.log(out, jacobians=jlog)
    for i in range(2):
        jac_analytic = jlog[0] @ jtest[i] if name == "compose" else jtest[i]
        torch.testing.assert_allclose(jacs_vmap[i], jac_analytic)
//  [python]
def check_jacrev_binary(group_fns, batch_size, name):
    assert name in ["compose", "transform", "untransform"]
    if not hasattr(torch, "vmap"):
        return

    test_fn = getattr(group_fns, name)

    fn_inputs = (
        (group_fns.rand(batch_size), torch.randn(batch_size, 3))
        if name in ["transform", "untransform"]
        else (group_fns.rand(batch_size), group_fns.rand(batch_size))
    )

    def f(t1, t2):
        op_out = test_fn(t1.unsqueeze(0), t2.unsqueeze(0))
        if name == "compose":
            op_out = group_fns.log(op_out)
        return op_out.squeeze(0)

    jacs_vmap = []
    for i in range(2):
        j = torch.vmap(torch.func.jacrev(f, i))(fn_inputs[0], fn_inputs[1])
        if fn_inputs[i].ndim == 3:  # group input
            j = group_fns.left_project(fn_inputs[i], j)
        jacs_vmap.append(j)

    jlog, jtest = [], []
    out = test_fn(fn_inputs[0], fn_inputs[1], jacobians=jtest)
    if name == "compose":
        group_fns.log(out, jacobians=jlog)
    for i in range(2):
        jac_analytic = jlog[0] @ jtest[i] if name == "compose" else jtest[i]
        torch.testing.assert_close(jacs_vmap[i], jac_analytic)
//  [python]
def test_one_hot_num_classes(self, targets):
        num_classes = targets.max() + 1
        composer_one_hot = _one_hot(targets, num_classes=num_classes)
        pytorch_one_hot = F.one_hot(targets, num_classes=num_classes)
        torch.testing.assert_allclose(composer_one_hot, pytorch_one_hot)
//  [python]
def test_one_hot_num_classes(self, targets):
        num_classes = targets.max() + 1
        composer_one_hot = _one_hot(targets, num_classes=num_classes)
        pytorch_one_hot = F.one_hot(targets, num_classes=num_classes)
        torch.testing.assert_close(composer_one_hot, pytorch_one_hot)
//  [python]
def check_binary_op_broadcasting(group_fns, op_name, group_size, bs1, bs2, dtype, rng):
    assert op_name in ["compose", "transform", "untransform"]
    g1 = group_fns.rand(*bs1, generator=rng, dtype=dtype)
    if op_name == "compose":
        t2 = group_fns.rand(*bs2, generator=rng, dtype=dtype)
        t2_size = group_size
    else:
        t2 = torch.randn(*bs2, 3, generator=rng, dtype=dtype)
        t2_size = (3,)

    # The following code does broadcasting manually, then we check that
    # manual broadcast output is the same as the automatic broadcasting
    broadcast_size = _get_broadcast_size(bs1, bs2)
    t1_expand_flat = _expand_flat(g1, broadcast_size, group_size)
    t2_expand_flat = _expand_flat(t2, broadcast_size, t2_size)

    fn = getattr(group_fns, op_name)
    jfn = getattr(group_fns, f"j{op_name}")
    out = fn(g1, t2)
    out_expand_flat = fn(t1_expand_flat, t2_expand_flat)
    t.assert_close(out, out_expand_flat.reshape(broadcast_size + t2_size))

    jout = jfn(g1, t2)[0]
    jout_expand_flat = jfn(t1_expand_flat, t2_expand_flat)[0]
    for j1, j2 in zip(jout, jout_expand_flat):
        t.assert_close(j1, j2.reshape(broadcast_size + j1.shape[-2:]))
//  [python]
def check_binary_op_broadcasting(group_fns, op_name, group_size, bs1, bs2, dtype, rng):
    assert op_name in ["compose", "transform", "untransform"]
    g1 = group_fns.rand(*bs1, generator=rng, dtype=dtype)
    if op_name == "compose":
        t2 = group_fns.rand(*bs2, generator=rng, dtype=dtype)
        t2_size = group_size
    else:
        t2 = torch.randn(*bs2, 3, generator=rng, dtype=dtype)
        t2_size = (3,)

    # The following code does broadcasting manually, then we check that
    # manual broadcast output is the same as the automatic broadcasting
    broadcast_size = _get_broadcast_size(bs1, bs2)
    t1_expand_flat = _expand_flat(g1, broadcast_size, group_size)
    t2_expand_flat = _expand_flat(t2, broadcast_size, t2_size)

    fn = getattr(group_fns, op_name)
    jfn = getattr(group_fns, f"j{op_name}")
    out = fn(g1, t2)
    out_expand_flat = fn(t1_expand_flat, t2_expand_flat)
    torch.testing.assert_close(out, out_expand_flat.reshape(broadcast_size + t2_size))

    jout = jfn(g1, t2)[0]
    jout_expand_flat = jfn(t1_expand_flat, t2_expand_flat)[0]
    for j1, j2 in zip(jout, jout_expand_flat):
        torch.testing.assert_close(j1, j2.reshape(broadcast_size + j1.shape[-2:]))
//  [python]
def test_ensure_targets_one_hot(tensors):
    input, targets_idx, targets_one_hot = tensors
    targets_one_hot_test = ensure_targets_one_hot(input, targets_idx)
    t.assert_allclose(targets_one_hot, targets_one_hot_test, check_stride=False)
//  [python]
def test_ensure_targets_one_hot(tensors):
    input, targets_idx, targets_one_hot = tensors
    targets_one_hot_test = ensure_targets_one_hot(input, targets_idx)
    torch.testing.assert_close(targets_one_hot, targets_one_hot_test, check_stride=False)
//  [python]
def test_gemm_fusion():
    M, N, K = 4096, 4096, 64
    BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 64
    A = torch.empty((M, K), dtype=torch.float16, device='cuda').normal_(mean=0.1, std=0.2)
    B = torch.empty((N, K), dtype=torch.float16, device='cuda').normal_(mean=0.1, std=0.2)
    C = torch.empty((N, K), dtype=torch.float16, device='cuda').normal_(mean=0.1, std=0.2)
    E = torch.empty((M, K), dtype=torch.float16, device='cuda')
    ref_out = torch.matmul(torch.matmul(A, B.T), C)
    num_warps = 4
    grid = (triton.cdiv(M, BLOCK_M), 1)
    gemm_fusion_kernel[grid](
        A, B, C, E, M, N, K,  #
        A.stride(0), A.stride(1),  #
        B.stride(0), B.stride(1),  #
        C.stride(0), C.stride(1),  #
        E.stride(0), E.stride(1),  #
        BLOCK_M, BLOCK_N, BLOCK_K,  #
        num_warps=num_warps)

    t.assert_allclose(ref_out, E, atol=1e-2, rtol=0)
//  [python]
def test_gemm_fusion():
    M, N, K = 4096, 4096, 64
    BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 64
    A = torch.empty((M, K), dtype=torch.float16, device='cuda').normal_(mean=0.1, std=0.2)
    B = torch.empty((N, K), dtype=torch.float16, device='cuda').normal_(mean=0.1, std=0.2)
    C = torch.empty((N, K), dtype=torch.float16, device='cuda').normal_(mean=0.1, std=0.2)
    E = torch.empty((M, K), dtype=torch.float16, device='cuda')
    ref_out = torch.matmul(torch.matmul(A, B.T), C)
    num_warps = 4
    grid = (triton.cdiv(M, BLOCK_M), 1)
    gemm_fusion_kernel[grid](
        A, B, C, E, M, N, K,  #
        A.stride(0), A.stride(1),  #
        B.stride(0), B.stride(1),  #
        C.stride(0), C.stride(1),  #
        E.stride(0), E.stride(1),  #
        BLOCK_M, BLOCK_N, BLOCK_K,  #
        num_warps=num_warps)

    torch.testing.assert_close(ref_out, E, atol=1e-2, rtol=0)
//  [python]
def test_one_hot_dim(self, targets, dim):
        composer_one_hot = _one_hot(targets, dim=dim)
        pytorch_one_hot = F.one_hot(targets)
        # Move class dim to specified dim
        pytorch_one_hot = torch.movedim(pytorch_one_hot, source=-1, destination=dim).contiguous()
        torch.testing.assert_allclose(composer_one_hot, pytorch_one_hot)
//  [python]
def test_one_hot_dim(self, targets, dim):
        composer_one_hot = _one_hot(targets, dim=dim)
        pytorch_one_hot = F.one_hot(targets)
        # Move class dim to specified dim
        pytorch_one_hot = torch.movedim(pytorch_one_hot, source=-1, destination=dim).contiguous()
        torch.testing.assert_close(composer_one_hot, pytorch_one_hot)
//  [python]
def test_incorrect_prediction(
        self,
        incorrect_input: torch.Tensor,
        target: torch.Tensor,
        squared_pred: bool,
        jaccard: bool,
        batch: bool,
        ignore_absent_classes: bool,
    ):
        dice_loss = DiceLoss(
            squared_pred=squared_pred,
            jaccard=jaccard,
            batch=batch,
            ignore_absent_classes=ignore_absent_classes,
        )
        loss = dice_loss(incorrect_input, target)
        torch.testing.assert_allclose(loss, torch.tensor(1.0))
//  [python]
def test_incorrect_prediction(
        self,
        incorrect_input: torch.Tensor,
        target: torch.Tensor,
        squared_pred: bool,
        jaccard: bool,
        batch: bool,
        ignore_absent_classes: bool,
    ):
        dice_loss = DiceLoss(
            squared_pred=squared_pred,
            jaccard=jaccard,
            batch=batch,
            ignore_absent_classes=ignore_absent_classes,
        )
        loss = dice_loss(incorrect_input, target)
        torch.testing.assert_close(loss, torch.tensor(1.0))
//  [python]
def test_soft_cross_entropy(self, tensors, use_weights, reduction):
        input, target_indices, target_onehot = tensors
        if use_weights:
            num_classes = target_onehot.shape[1]
            weights = torch.rand(size=[num_classes])
        else:
            weights = None

        loss_indices = soft_cross_entropy(input, target_indices, weight=weights, reduction=reduction, ignore_index=-1)
        loss_onehot = soft_cross_entropy(input, target_onehot, weight=weights, reduction=reduction)
        loss_reference = F.cross_entropy(input, target_indices, weight=weights, reduction=reduction, ignore_index=-1)

        self.assertEqual(loss_indices, loss_onehot)
        self.assertEqual(loss_indices, loss_reference)
//  [python]
def test_soft_cross_entropy(self, tensors, use_weights, reduction):
        input, target_indices, target_onehot = tensors
        if use_weights:
            num_classes = target_onehot.shape[1]
            weights = torch.rand(size=[num_classes])
        else:
            weights = None

        loss_indices = soft_cross_entropy(input, target_indices, weight=weights, reduction=reduction, ignore_index=-1)
        loss_onehot = soft_cross_entropy(input, target_onehot, weight=weights, reduction=reduction)
        loss_reference = F.cross_entropy(input, target_indices, weight=weights, reduction=reduction, ignore_index=-1)

        torch.testing.assert_close(loss_indices, loss_onehot)
        torch.testing.assert_close(loss_indices, loss_reference)
//  [python]
def _test_against_cpu(self, inputs, reference, eps, dtype, is_causal=None):
        q, k, v, rpb, kernel_size, dilation = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        attn_ref, out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            attn = na2d_qk(
                q_,
                k_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )
            attn = attn.softmax(dim=-1)
            out = na2d_av(
                attn,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_allclose(attn.float(), attn_ref, atol=eps, rtol=0)
            torch.testing.assert_allclose(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_cpu(self, inputs, reference, eps, dtype, is_causal=None):
        q, k, v, rpb, kernel_size, dilation = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        attn_ref, out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            attn = na2d_qk(
                q_,
                k_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )
            attn = attn.softmax(dim=-1)
            out = na2d_av(
                attn,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_close(attn.float(), attn_ref, atol=eps, rtol=0)
            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def test_setting_gains_from_cfg_dict(self):
        """Test that gains are loaded from the configuration dictionary correctly.

        Note: We purposefully give one argument as int and other as float to check that it is handled correctly.
        """
        # Create articulation
        robot_cfg = ArticulationCfg(
            prim_path="/World/Robot",
            spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Humanoid/humanoid_instanceable.usd"),
            init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)),
            actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness={".*": 10}, damping={".*": 2.0})},
        )
        robot = Articulation(cfg=robot_cfg)

        # Play sim
        self.sim.reset()

        # Expected gains
        expected_stiffness = torch.full((robot.num_instances, robot.num_joints), 10.0, device=robot.device)
        expected_damping = torch.full_like(expected_stiffness, 2.0)

        # Check that gains are loaded from USD file
        torch.testing.assert_allclose(robot.actuators["body"].stiffness, expected_stiffness)
        torch.testing.assert_allclose(robot.actuators["body"].damping, expected_damping)
//  [python]
def test_setting_gains_from_cfg_dict(self):
        """Test that gains are loaded from the configuration dictionary correctly.

        Note: We purposefully give one argument as int and other as float to check that it is handled correctly.
        """
        # Create articulation
        robot_cfg = ArticulationCfg(
            prim_path="/World/Robot",
            spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Humanoid/humanoid_instanceable.usd"),
            init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)),
            actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness={".*": 10}, damping={".*": 2.0})},
        )
        robot = Articulation(cfg=robot_cfg)

        # Play sim
        self.sim.reset()

        # Expected gains
        expected_stiffness = torch.full((robot.num_instances, robot.num_joints), 10.0, device=robot.device)
        expected_damping = torch.full_like(expected_stiffness, 2.0)

        # Check that gains are loaded from USD file
        torch.testing.assert_close(robot.actuators["body"].stiffness, expected_stiffness)
        torch.testing.assert_close(robot.actuators["body"].damping, expected_damping)
//  [python]
def test_setting_gains_from_cfg(self):
        """Test that gains are loaded from the configuration correctly.

        Note: We purposefully give one argument as int and other as float to check that it is handled correctly.
        """
        # Create articulation
        robot_cfg = ArticulationCfg(
            prim_path="/World/Robot",
            spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Humanoid/humanoid_instanceable.usd"),
            init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)),
            actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10, damping=2.0)},
        )
        robot = Articulation(cfg=robot_cfg)

        # Play sim
        self.sim.reset()

        # Expected gains
        expected_stiffness = torch.full((robot.num_instances, robot.num_joints), 10.0, device=robot.device)
        expected_damping = torch.full_like(expected_stiffness, 2.0)

        # Check that gains are loaded from USD file
        torch.testing.assert_allclose(robot.actuators["body"].stiffness, expected_stiffness)
        torch.testing.assert_allclose(robot.actuators["body"].damping, expected_damping)
//  [python]
def test_setting_gains_from_cfg(self):
        """Test that gains are loaded from the configuration correctly.

        Note: We purposefully give one argument as int and other as float to check that it is handled correctly.
        """
        # Create articulation
        robot_cfg = ArticulationCfg(
            prim_path="/World/Robot",
            spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Humanoid/humanoid_instanceable.usd"),
            init_state=ArticulationCfg.InitialStateCfg(pos=(0.0, 0.0, 1.34)),
            actuators={"body": ImplicitActuatorCfg(joint_names_expr=[".*"], stiffness=10, damping=2.0)},
        )
        robot = Articulation(cfg=robot_cfg)

        # Play sim
        self.sim.reset()

        # Expected gains
        expected_stiffness = torch.full((robot.num_instances, robot.num_joints), 10.0, device=robot.device)
        expected_damping = torch.full_like(expected_stiffness, 2.0)

        # Check that gains are loaded from USD file
        torch.testing.assert_close(robot.actuators["body"].stiffness, expected_stiffness)
        torch.testing.assert_close(robot.actuators["body"].damping, expected_damping)
//  [python]
def test_scale_coords():
    box_tensor = torch.tensor(
        [
            [0.0, 0.0, 100.0, 100.0],
            [0.0, 0.0, 0.0, 0.0],
            [10.0, 15.0, 30.0, 35.0],
            [20.0, 35.0, 90.0, 95.0],
        ],
        dtype=torch.float,
    )
    exp_coords = torch.tensor(
        [
            [0.0, 0.0, 108.05, 111.25],
            [0.0, 0.0, 0.0, 0.0],
            [7.9250, 16.6875, 30.1750, 38.9375],
            [19.05, 38.9375, 96.9250, 105.6875],
        ],
        dtype=torch.float,
    )

    box_coords_scaled = scale_coords((160, 128), box_tensor, (178, 136))
    assert tuple(box_coords_scaled.shape) == (4, 4)
    torch.testing.assert_allclose(box_coords_scaled, exp_coords)
//  [python]
def test_scale_coords():
    box_tensor = torch.tensor(
        [
            [0.0, 0.0, 100.0, 100.0],
            [0.0, 0.0, 0.0, 0.0],
            [10.0, 15.0, 30.0, 35.0],
            [20.0, 35.0, 90.0, 95.0],
        ],
        dtype=torch.float,
    )
    exp_coords = torch.tensor(
        [
            [0.0, 0.0, 108.05, 111.25],
            [0.0, 0.0, 0.0, 0.0],
            [7.9250, 16.6875, 30.1750, 38.9375],
            [19.05, 38.9375, 96.9250, 105.6875],
        ],
        dtype=torch.float,
    )

    box_coords_scaled = scale_coords((160, 128), box_tensor, (178, 136))
    assert tuple(box_coords_scaled.shape) == (4, 4)
    torch.testing.assert_close(box_coords_scaled, exp_coords)
//  [python]
def test_triton_contiguous_torch_op():

    def call_triton_contiguous(x, memory_format=torch.contiguous_format):
        return sfast_triton.contiguous(x, memory_format=memory_format)

    a = torch.ones(1, 4, 256, 512).cuda().permute(0, 1, 3, 2)

    out = call_triton_contiguous(a)
    assert out.is_contiguous()
    torch.testing.assert_allclose(out, a)

    out = call_triton_contiguous(a, memory_format=torch.channels_last)
    assert out.is_contiguous(memory_format=torch.channels_last)
    torch.testing.assert_allclose(out, a)

    traced_triton_contiguous = torch.jit.trace(call_triton_contiguous, (a,))
    out = traced_triton_contiguous(a)
    assert out.is_contiguous()
    torch.testing.assert_allclose(out, a)
    logger.info(f'Graph: {traced_triton_contiguous.graph}')
//  [python]
def test_triton_contiguous_torch_op():

    def call_triton_contiguous(x, memory_format=torch.contiguous_format):
        return sfast_triton.contiguous(x, memory_format=memory_format)

    a = torch.ones(1, 4, 256, 512).cuda().permute(0, 1, 3, 2)

    out = call_triton_contiguous(a)
    assert out.is_contiguous()
    torch.testing.assert_close(out, a)

    out = call_triton_contiguous(a, memory_format=torch.channels_last)
    assert out.is_contiguous(memory_format=torch.channels_last)
    torch.testing.assert_close(out, a)

    traced_triton_contiguous = torch.jit.trace(call_triton_contiguous, (a,))
    out = traced_triton_contiguous(a)
    assert out.is_contiguous()
    torch.testing.assert_close(out, a)
    logger.info(f'Graph: {traced_triton_contiguous.graph}')
//  [python]
def _test_transpose_static_shape(
        self,
        input_shape: Sequence[int],
        dim0: int,
        dim1: int,
        dtype: str = "float16",
        test_name: str = "transpose_static_shape",
    ) -> None:
        X = Tensor(
            shape=input_shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )
        op = ops.transpose()
        Y = op(X, dim0, dim1)
        Y._attrs["is_output"] = True
        Y._attrs["name"] = "output"
        target = detect_target()
        module = compile_model(Y, target, "./tmp", f"{test_name}_{self._test_id}")
        self._test_id += 1

        X_pt = get_random_torch_tensor(input_shape, dtype=dtype)
        Y_pt = torch.transpose(X_pt, dim0, dim1).contiguous()

        y = torch.empty_like(Y_pt)
        module.run_with_tensors([X_pt], [y])

        torch.testing.assert_allclose(y, Y_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_transpose_static_shape(
        self,
        input_shape: Sequence[int],
        dim0: int,
        dim1: int,
        dtype: str = "float16",
        test_name: str = "transpose_static_shape",
    ) -> None:
        X = Tensor(
            shape=input_shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )
        op = ops.transpose()
        Y = op(X, dim0, dim1)
        Y._attrs["is_output"] = True
        Y._attrs["name"] = "output"
        target = detect_target()
        module = compile_model(Y, target, "./tmp", f"{test_name}_{self._test_id}")
        self._test_id += 1

        X_pt = get_random_torch_tensor(input_shape, dtype=dtype)
        Y_pt = torch.transpose(X_pt, dim0, dim1).contiguous()

        y = torch.empty_like(Y_pt)
        module.run_with_tensors([X_pt], [y])

        torch.testing.assert_close(y, Y_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7), (3, 2)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        query = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, 16),
                torch.randn(2, 8, 16, 18, 32),
                torch.randn(4, 1, 32, 20, 16),
            ],
            **kwargs,
        )
        key = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, 16),
                torch.randn(2, 8, 16, 18, 32),
                torch.randn(4, 1, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na2d_qk(query, key, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for q, k in zip(query, key):
            out_ref.append(na2d_qk(q, k, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            additional_keys = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na2d_qk(
                query,
                key,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_keys=additional_keys,
            )
            out_ref = []
            for q, k, k_add in zip(query, key, additional_keys):
                out_ref.append(
                    na2d_qk(
                        q,
                        k,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_keys=k_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def _test_nested_qk_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7), (3, 2)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        query = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, 16),
                torch.randn(2, 8, 16, 18, 32),
                torch.randn(4, 1, 32, 20, 16),
            ],
            **kwargs,
        )
        key = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, 16),
                torch.randn(2, 8, 16, 18, 32),
                torch.randn(4, 1, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na2d_qk(query, key, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for q, k in zip(query, key):
            out_ref.append(na2d_qk(q, k, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            additional_keys = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na2d_qk(
                query,
                key,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_keys=additional_keys,
            )
            out_ref = []
            for q, k, k_add in zip(query, key, additional_keys):
                out_ref.append(
                    na2d_qk(
                        q,
                        k,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_keys=k_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def assert_allclose(
    actual: Any,
    expected: Any,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    equal_nan: bool = True,
    msg: str = "",
) -> None:
    if not isinstance(actual, torch.Tensor):
        actual = torch.tensor(actual)
    if not isinstance(expected, torch.Tensor):
        expected = torch.tensor(expected, dtype=actual.dtype)

    if rtol is None and atol is None:
        rtol, atol = _get_default_rtol_and_atol(actual, expected)

    torch.testing.assert_allclose(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        equal_nan=equal_nan,
        check_device=True,
        check_dtype=False,
        check_stride=False,
        msg=msg or None,
    )
//  [python]
def assert_allclose(
    actual: Any,
    expected: Any,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    equal_nan: bool = True,
    msg: str = "",
) -> None:
    if not isinstance(actual, torch.Tensor):
        actual = torch.tensor(actual)
    if not isinstance(expected, torch.Tensor):
        expected = torch.tensor(expected, dtype=actual.dtype)

    if rtol is None and atol is None:
        rtol, atol = _get_default_rtol_and_atol(actual, expected)

    torch.testing.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        equal_nan=equal_nan,
        check_device=True,
        check_dtype=False,
        check_stride=False,
        msg=msg or None,
    )
//  [python]
def validate_ema(model, original_model, ema_model, smoothing):
    model_params, model_buffers = dict(model.named_parameters()), dict(model.named_buffers())
    original_params, original_buffers = dict(original_model.named_parameters()), dict(original_model.named_buffers())
    ema_params, ema_buffers = dict(ema_model.named_parameters()), dict(ema_model.named_buffers())

    for name, param in model_params.items():
        new_param = (original_params[name] * smoothing + (1. - smoothing) * param)
        torch.testing.assert_allclose(ema_params[name].data, new_param)

    for name, buffer in model_buffers.items():
        new_buffer = (original_buffers[name] * smoothing + (1. - smoothing) * buffer).type(ema_buffers[name].data.dtype)
        torch.testing.assert_allclose(ema_buffers[name].data, new_buffer)
//  [python]
def validate_ema(model, original_model, ema_model, smoothing):
    model_params, model_buffers = dict(model.named_parameters()), dict(model.named_buffers())
    original_params, original_buffers = dict(original_model.named_parameters()), dict(original_model.named_buffers())
    ema_params, ema_buffers = dict(ema_model.named_parameters()), dict(ema_model.named_buffers())

    for name, param in model_params.items():
        new_param = (original_params[name] * smoothing + (1. - smoothing) * param)
        torch.testing.assert_close(ema_params[name].data, new_param)

    for name, buffer in model_buffers.items():
        new_buffer = (original_buffers[name] * smoothing + (1. - smoothing) * buffer).type(ema_buffers[name].data.dtype)
        torch.testing.assert_close(ema_buffers[name].data, new_buffer)
//  [python]
def get_default_comparer(
    rtol: float = 1e-04,
    atol: float = 0,
    equal_nan: bool = True,
) -> _CompareFn:
    """Creates default comparer function.

    The created function will compare the outputs by using
    `torch.testing.assert_allclose` with specified options.

    Args:
        rtol (float): Relative tolerance.
        atol (float): Absolute tolerance.
        equal_nan (bool): If ``True``, NaNs will be ignored.
    """

    def compare_fn(
        backend1: str, backend2: str, name: str, val1: Any, val2: Any
    ) -> None:
        # TODO select the device where
        # the tensors will be compared?
        if isinstance(val1, torch.Tensor):
            val1 = val1.cpu().detach()
        if isinstance(val2, torch.Tensor):
            val2 = val2.cpu().detach()

        torch.testing.assert_allclose(
            val1, val2, rtol=rtol, atol=atol, equal_nan=equal_nan
        )

    return compare_fn
//  [python]
def get_default_comparer(
    rtol: float = 1e-04,
    atol: float = 0,
    equal_nan: bool = True,
) -> _CompareFn:
    """Creates default comparer function.

    The created function will compare the outputs by using
    `torch.testing.assert_allclose` with specified options.

    Args:
        rtol (float): Relative tolerance.
        atol (float): Absolute tolerance.
        equal_nan (bool): If ``True``, NaNs will be ignored.
    """

    def compare_fn(
        backend1: str, backend2: str, name: str, val1: Any, val2: Any
    ) -> None:
        # TODO select the device where
        # the tensors will be compared?
        if isinstance(val1, torch.Tensor):
            val1 = val1.cpu().detach()
        if isinstance(val2, torch.Tensor):
            val2 = val2.cpu().detach()

        torch.testing.assert_close(
            val1, val2, rtol=rtol, atol=atol, equal_nan=equal_nan
        )

    return compare_fn
//  [python]
def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7), (3, 2)
        num_na_weights = get_num_na_weights(kernel_size)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        attn = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, num_na_weights),
                torch.randn(2, 8, 16, 18, num_na_weights),
                torch.randn(4, 1, 32, 20, num_na_weights),
            ],
            **kwargs,
        )
        value = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, 16),
                torch.randn(2, 8, 16, 18, 32),
                torch.randn(4, 1, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na2d_av(attn, value, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for a, v in zip(attn, value):
            out_ref.append(na2d_av(a, v, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            attn = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, 14, 16, num_na_weights + n_extra_tokens),
                    torch.randn(2, 8, 16, 18, num_na_weights + n_extra_tokens),
                    torch.randn(4, 1, 32, 20, num_na_weights + n_extra_tokens),
                ],
                **kwargs,
            )
            additional_values = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na2d_av(
                attn,
                value,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_values=additional_values,
            )
            out_ref = []
            for a, v, v_add in zip(attn, value, additional_values):
                out_ref.append(
                    na2d_av(
                        a,
                        v,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_values=v_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_allclose(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def _test_nested_av_forward(self, dtype, device, test_additional_tokens=True):
        kernel_size, dilation = (3, 7), (3, 2)
        num_na_weights = get_num_na_weights(kernel_size)
        kwargs = {"dtype": dtype, "device": device, "requires_grad": False}
        attn = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, num_na_weights),
                torch.randn(2, 8, 16, 18, num_na_weights),
                torch.randn(4, 1, 32, 20, num_na_weights),
            ],
            **kwargs,
        )
        value = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, 14, 16, 16),
                torch.randn(2, 8, 16, 18, 32),
                torch.randn(4, 1, 32, 20, 16),
            ],
            **kwargs,
        )
        out_nested = na2d_av(attn, value, kernel_size=kernel_size, dilation=dilation)
        out_ref = []
        for a, v in zip(attn, value):
            out_ref.append(na2d_av(a, v, kernel_size=kernel_size, dilation=dilation))

        for o, o_ref in zip(out_nested, out_ref):
            torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)

        if test_additional_tokens:
            n_extra_tokens = 9
            attn = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, 14, 16, num_na_weights + n_extra_tokens),
                    torch.randn(2, 8, 16, 18, num_na_weights + n_extra_tokens),
                    torch.randn(4, 1, 32, 20, num_na_weights + n_extra_tokens),
                ],
                **kwargs,
            )
            additional_values = torch.nested.nested_tensor(
                [
                    torch.randn(1, 2, n_extra_tokens, 16),
                    torch.randn(2, 8, n_extra_tokens, 32),
                    torch.randn(4, 1, n_extra_tokens, 16),
                ],
                **kwargs,
            )
            out_nested = na2d_av(
                attn,
                value,
                kernel_size=kernel_size,
                dilation=dilation,
                additional_values=additional_values,
            )
            out_ref = []
            for a, v, v_add in zip(attn, value, additional_values):
                out_ref.append(
                    na2d_av(
                        a,
                        v,
                        kernel_size=kernel_size,
                        dilation=dilation,
                        additional_values=v_add,
                    )
                )

            for o, o_ref in zip(out_nested, out_ref):
                torch.testing.assert_close(o, o_ref, atol=1e-6, rtol=0)
//  [python]
def validate_model(model1, model2):
    model1_params, model1_buffers = dict(model1.named_parameters()), dict(model1.named_buffers())
    model2_params, model2_buffers = dict(model2.named_parameters()), dict(model2.named_buffers())

    for name in model1_params.keys():
        t.assert_close(model1_params[name].data, model2_params[name].data)

    for name in model1_buffers.keys():
        t.assert_close(model1_buffers[name].data, model2_buffers[name].data)
//  [python]
def validate_model(model1, model2):
    model1_params, model1_buffers = dict(model1.named_parameters()), dict(model1.named_buffers())
    model2_params, model2_buffers = dict(model2.named_parameters()), dict(model2.named_buffers())

    for name in model1_params.keys():
        torch.testing.assert_close(model1_params[name].data, model2_params[name].data)

    for name in model1_buffers.keys():
        torch.testing.assert_close(model1_buffers[name].data, model2_buffers[name].data)
//  [python]
def _validate_input_feature(self, features: FeaturesType) -> None:
        """
        Validates that the proper feature is provided.
        Raises an exception if it is not.
        :param features: The feature provided to the model.
        """
        if "MockFeature" not in features:
            raise ValueError(f"MockFeature not in provided features. Available keys: {sorted(list(features.keys()))}")

        if len(features) != 1:
            raise ValueError(f"Expected a single feature. Instead got {len(features)}: {sorted(list(features.keys()))}")

        mock_feature = features["MockFeature"]
        if not isinstance(mock_feature, MockFeature):
            raise ValueError(f"Expected feature of type MockFeature, but got {type(mock_feature)}")

        mock_feature_data = mock_feature.data
        t.assert_close(mock_feature_data, self.expected_forward_tensor)
//  [python]
def _validate_input_feature(self, features: FeaturesType) -> None:
        """
        Validates that the proper feature is provided.
        Raises an exception if it is not.
        :param features: The feature provided to the model.
        """
        if "MockFeature" not in features:
            raise ValueError(f"MockFeature not in provided features. Available keys: {sorted(list(features.keys()))}")

        if len(features) != 1:
            raise ValueError(f"Expected a single feature. Instead got {len(features)}: {sorted(list(features.keys()))}")

        mock_feature = features["MockFeature"]
        if not isinstance(mock_feature, MockFeature):
            raise ValueError(f"Expected feature of type MockFeature, but got {type(mock_feature)}")

        mock_feature_data = mock_feature.data
        torch.testing.assert_close(mock_feature_data, self.expected_forward_tensor)
//  [python]
def test_make_callable_from_trace(executor, device: str, dtype: dtypes.dtype):
    def foo(a, b):
        return a + b

    a = make_tensor((2, 2), device=device, dtype=torch.float32)
    b = make_tensor((2, 2), device=device, dtype=torch.float32)
    traced_foo = thunder.trace(inline_trace=False)(foo, a, b)
    assert len(traced_foo.bound_symbols) == 4
    assert traced_foo.bound_symbols[-2].sym.name == "add"

    # Ensure that the trace can be fused and executed
    fusion = executor.make_callable(traced_foo)
    actual = fusion(a, b)
    t.assert_close(actual, a + b)
//  [python]
def test_make_callable_from_trace(executor, device: str, dtype: dtypes.dtype):
    def foo(a, b):
        return a + b

    a = make_tensor((2, 2), device=device, dtype=torch.float32)
    b = make_tensor((2, 2), device=device, dtype=torch.float32)
    traced_foo = thunder.trace(inline_trace=False)(foo, a, b)
    assert len(traced_foo.bound_symbols) == 4
    assert traced_foo.bound_symbols[-2].sym.name == "add"

    # Ensure that the trace can be fused and executed
    fusion = executor.make_callable(traced_foo)
    actual = fusion(a, b)
    torch.testing.assert_close(actual, a + b)
//  [python]
def test_integer_return(executor, device, _):
    if executor == nvFuserExecutor:
        pytest.xfail("nvFuser does not support only scalar outputs")

    def foo(a, b):
        return clang.add(a, b)

    traced_foo = executor.make_callable(foo)

    thunder_result = traced_foo(3, 4)
    python_result = 3 + 4
    t.assert_close(thunder_result, python_result)
//  [python]
def test_integer_return(executor, device, _):
    if executor == nvFuserExecutor:
        pytest.xfail("nvFuser does not support only scalar outputs")

    def foo(a, b):
        return clang.add(a, b)

    traced_foo = executor.make_callable(foo)

    thunder_result = traced_foo(3, 4)
    python_result = 3 + 4
    assert_close(thunder_result, python_result)
//  [python]
def test_varargs(executor, device, dtype):
    def foo(*args):
        return reduce(operator.add, args)

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    packed = (a, a, a, a, a)

    thunder_result = traced_foo(*packed)
    torch_result = foo(*packed)

    t.assert_close(thunder_result, torch_result)
//  [python]
def test_varargs(executor, device, dtype):
    def foo(*args):
        return reduce(operator.add, args)

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    packed = (a, a, a, a, a)

    thunder_result = traced_foo(*packed)
    torch_result = foo(*packed)

    assert_close(thunder_result, torch_result)
//  [python]
def test_kwargs(executor, device, dtype):
    def foo(**kwargs):
        return kwargs["a"] + kwargs["b"]

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2,), device=device, dtype=tdtype)

    thunder_result = traced_foo(a=a, b=b)
    torch_result = foo(a=a, b=b)

    t.assert_close(thunder_result, torch_result)
//  [python]
def test_kwargs(executor, device, dtype):
    def foo(**kwargs):
        return kwargs["a"] + kwargs["b"]

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2,), device=device, dtype=tdtype)

    thunder_result = traced_foo(a=a, b=b)
    torch_result = foo(a=a, b=b)

    assert_close(thunder_result, torch_result)
//  [python]
def test_no_return(executor, device, dtype):
    def foo(a, b):
        c = a + b  # noqa
        pass

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2, 2, 2), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b=b)
    torch_result = foo(a, b)

    t.assert_close(thunder_result, torch_result)
//  [python]
def test_no_return(executor, device, dtype):
    def foo(a, b):
        c = a + b  # noqa
        pass

    traced_foo = executor.make_callable(foo)
    tdtype = ltorch.to_torch_dtype(dtype)

    a = make_tensor((2,), device=device, dtype=tdtype)
    b = make_tensor((2, 2, 2), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b=b)
    torch_result = foo(a, b)

    assert_close(thunder_result, torch_result)
//  [python]
def test_no_input(executor, device, dtype):
    def foo():
        return 3, ()

    traced_foo = executor.make_callable(foo)

    thunder_result = traced_foo()
    torch_result = foo()

    t.assert_close(thunder_result, torch_result)
//  [python]
def test_no_input(executor, device, dtype):
    def foo():
        return 3, ()

    traced_foo = executor.make_callable(foo)

    thunder_result = traced_foo()
    torch_result = foo()

    assert_close(thunder_result, torch_result)
//  [python]
def test_integer_isinstance_mimicry(executor, device: str, dtype: dtypes.dtype):
    # isinstance() works as expected
    def foo(a, b, c):
        if isinstance(a, int):
            return clang.add(a, b)

        return clang.add(b, c)

    traced_foo = executor.make_callable(foo)

    tdtype = ltorch.to_torch_dtype(dtype)
    a = make_tensor((2, 1), device=device, dtype=tdtype)
    b = make_tensor((2, 2), device=device, dtype=tdtype)
    c = make_tensor((1, 2), device=device, dtype=tdtype)

    thunder_result = traced_foo(a, b, c)
    torch_result = b + c
    t.assert_close(thunder_result, torch_result)

    thunder_result = traced_foo(2, b, c)
    torch_result = 2 + b
    t.assert_close(thunder_result, torch_result)

    # type() doesn't work (it returns the actual type)
    def bar(a, b, c):
        if type(a) is int:
            return clang.add(a, b)

        return clang.add(b, c)

    traced_bar = executor.make_callable(bar)

    try:
        thunder_result = traced_bar(a, b, c)
        torch_result = b + c
        t.assert_close(thunder_result, torch_result)
        pytest.fail()
    except BaseException:
        pass

    try:
        thunder_result = traced_bar(2, b, c)
        torch_result = 2 + b
        t.assert_close(thunder_result, torch_result)
        pytest.fail()
    except BaseException:
        pass