jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def woq_assert_near_eq(ref, act, wTypeId):
    # match the scale in cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.cpp
    if wTypeId == 1:
        bits_in_type = 8
    else:
        bits_in_type = 4
    quant_range_scale = 1.0 / float(1 << (bits_in_type - 1))

    max_val = torch.max(abs(ref)).item()
    atol = (max_val * quant_range_scale) * 1.5  # allow for rounding
    torch.testing.assert_close(ref.cpu(), act.cpu(), atol=atol, rtol=1e-7)
//  [python]
def test_identity_reshape_in_out_conflict(self):
        """
        If x is an input and y is an output tensor, then we can only eliminate
        one view op in the following example:

          y = reshape(reshape(x, y_shape), x_original_shape)
        """
        dtype = "float"
        x0_shape = [2, 4, 8]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])

        y = ops.reshape()(x1, x0_shape)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(x0_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_identity_reshape_in_out_conflict",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 2)
        self.assertTrue(graph_has_op(result_graph, "reshape"))
        torch.testing.assert_allclose(x0_pt, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def test_identity_reshape_in_out_conflict(self):
        """
        If x is an input and y is an output tensor, then we can only eliminate
        one view op in the following example:

          y = reshape(reshape(x, y_shape), x_original_shape)
        """
        dtype = "float"
        x0_shape = [2, 4, 8]

        x0 = Tensor(
            shape=x0_shape,
            dtype=dtype,
            name="x0",
            is_input=True,
        )
        x1 = ops.reshape()(x0, [8, 8])

        y = ops.reshape()(x1, x0_shape)
        y._attrs["name"] = "y"
        y._attrs["is_output"] = True

        x0_pt = get_random_torch_tensor(x0_shape, dtype)
        y_pt = get_torch_empty_tensor(x0_shape, dtype)

        target = detect_target()
        module = compile_model(
            y,
            target,
            "./tmp",
            "test_identity_reshape_in_out_conflict",
        )
        result_graph = module.debug_sorted_graph
        module.run_with_tensors({"x0": x0_pt}, {"y": y_pt})

        self.assertEqual(len(result_graph), 2)
        self.assertTrue(graph_has_op(result_graph, "reshape"))
        torch.testing.assert_close(x0_pt, y_pt, atol=5e-2, rtol=5e-2)
//  [python]
def test_quantize_op(self, nrows: int, ncols: int) -> None:
        ebits = 8
        mbits = 7
        bias = 127
        max_pos = (1 << ((1 << ebits) - 2 - bias)) * (2 - 2 ** (-mbits))
        min_pos = 2 ** (1 - bias - mbits)
        bounding_box_size = 16
        print("MSFP parameters", bounding_box_size, ebits, mbits, bias)
        input_data = torch.rand(nrows, ncols).float()
        quantized_data = torch.ops.fbgemm.FloatToMSFPQuantized(
            input_data.cuda(),
            bounding_box_size,
            ebits,
            mbits,
            bias,
            min_pos,
            max_pos,
        )
        dequantized_data = torch.ops.fbgemm.MSFPQuantizedToFloat(
            quantized_data.cuda(), ebits, mbits, bias
        )
        torch.testing.assert_allclose(dequantized_data.cpu(), input_data, rtol=1, atol=0)
//  [python]
def test_quantize_op(self, nrows: int, ncols: int) -> None:
        ebits = 8
        mbits = 7
        bias = 127
        max_pos = (1 << ((1 << ebits) - 2 - bias)) * (2 - 2 ** (-mbits))
        min_pos = 2 ** (1 - bias - mbits)
        bounding_box_size = 16
        print("MSFP parameters", bounding_box_size, ebits, mbits, bias)
        input_data = torch.rand(nrows, ncols).float()
        quantized_data = torch.ops.fbgemm.FloatToMSFPQuantized(
            input_data.cuda(),
            bounding_box_size,
            ebits,
            mbits,
            bias,
            min_pos,
            max_pos,
        )
        dequantized_data = torch.ops.fbgemm.MSFPQuantizedToFloat(
            quantized_data.cuda(), ebits, mbits, bias
        )
        torch.testing.assert_close(dequantized_data.cpu(), input_data, rtol=1, atol=0)
//  [python]
def run_worker(args):
    mod = ExampleCode()
    mod.to(args.device)

    x = torch.randn(batch_size, d_hid, device=args.device)
    y = torch.randn(batch_size, d_hid, device=args.device)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
        example_kwargs={"y": y},
    )

    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Attach to a schedule
    schedule = ScheduleGPipe(stage, args.chunks)

    # Run
    if args.rank == 0:
        schedule.step(x, y=y)
    else:
        out = schedule.step()

    dist.barrier()
    print(f"Rank {args.rank} completes")

    # Last rank checks result
    if args.rank == args.world_size - 1:
        ref_out = mod(x, y=y)
        t.assert_close(out, ref_out)
        print(
            f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref_out)}"
        )

    # Test qualname mapping
    submod_keys = stage.submod.state_dict().keys()
    print(f"Rank {args.rank} state dict keys: {submod_keys}")
    # Confirm keys are consistent with original model
    old_keys = mod.state_dict().keys()
    assert all(k in old_keys for k in submod_keys)
    print(f"Qualname test passed")
//  [python]
def run_worker(args):
    mod = ExampleCode()
    mod.to(args.device)

    x = torch.randn(batch_size, d_hid, device=args.device)
    y = torch.randn(batch_size, d_hid, device=args.device)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x,),
        example_kwargs={"y": y},
    )

    stage = PipelineStage(
        pipe,
        args.rank,
        device=args.device,
    )

    # Attach to a schedule
    schedule = ScheduleGPipe(stage, args.chunks)

    # Run
    if args.rank == 0:
        schedule.step(x, y=y)
    else:
        out = schedule.step()

    dist.barrier()
    print(f"Rank {args.rank} completes")

    # Last rank checks result
    if args.rank == args.world_size - 1:
        ref_out = mod(x, y=y)
        torch.testing.assert_close(out, ref_out)
        print(
            f"equivalence test passed {torch.sum(out)} ref {torch.sum(ref_out)}"
        )

    # Test qualname mapping
    submod_keys = stage.submod.state_dict().keys()
    print(f"Rank {args.rank} state dict keys: {submod_keys}")
    # Confirm keys are consistent with original model
    old_keys = mod.state_dict().keys()
    assert all(k in old_keys for k in submod_keys)
    print(f"Qualname test passed")
//  [python]
def test_bfloat16_llama_init(lit_llama, orig_llama) -> None:
    from lit_llama.utils import EmptyInitOnDevice

    block_size = 64
    vocab_size = 32000
    n_layer = 16
    n_head = 16
    n_embd = 32

    llama_config = lit_llama.LLaMAConfig(
        block_size=block_size, vocab_size=vocab_size, n_layer=n_layer, n_head=n_head, n_embd=n_embd
    )
    llama_model = lit_llama.LLaMA(llama_config)
    llama_model.apply(llama_model._init_weights)

    batch_size = 3

    token_sample = torch.randint(0, vocab_size, size=(batch_size, block_size), dtype=torch.int64)

    expected = llama_model(token_sample)

    with EmptyInitOnDevice(device="cuda", dtype=torch.bfloat16):
        llama_model2 = lit_llama.LLaMA(llama_config)
    llama_model2.load_state_dict(llama_model.state_dict(keep_vars=True))

    out = llama_model2(token_sample.cuda()).float().cpu()
    torch.allclose(out, expected, atol=5e-3, rtol=1e-3)
//  [python]
def test_bfloat16_llama_init(lit_llama, orig_llama) -> None:
    from lit_llama.utils import EmptyInitOnDevice

    block_size = 64
    vocab_size = 32000
    n_layer = 16
    n_head = 16
    n_embd = 32

    llama_config = lit_llama.LLaMAConfig(
        block_size=block_size, vocab_size=vocab_size, n_layer=n_layer, n_head=n_head, n_embd=n_embd
    )
    llama_model = lit_llama.LLaMA(llama_config)
    llama_model.apply(llama_model._init_weights)

    batch_size = 3

    token_sample = torch.randint(0, vocab_size, size=(batch_size, block_size), dtype=torch.int64)

    expected = llama_model(token_sample)

    with EmptyInitOnDevice(device="cuda", dtype=torch.bfloat16):
        llama_model2 = lit_llama.LLaMA(llama_config)
    llama_model2.load_state_dict(llama_model.state_dict(keep_vars=True))

    out = llama_model2(token_sample.cuda()).float().cpu()
    torch.testing.assert_close(out, expected, atol=5e-3, rtol=1e-3)
//  [python]
def _eval_models(self, blade_model: WrapperModel, eager_model: PreTrainedModel) -> None:
        inputs = blade_model.info.example_inputs
        blade_outputs = blade_model(**inputs)
        eager_outputs = eager_model(**inputs)
        self.assertTrue(blade_outputs)
        for x in blade_outputs:
            torch.allclose(
                blade_outputs[x].float(), eager_outputs[x].float(), rtol=0.001, atol=0.001)
//  [python]
def _eval_models(self, blade_model: WrapperModel, eager_model: PreTrainedModel) -> None:
        inputs = blade_model.info.example_inputs
        blade_outputs = blade_model(**inputs)
        eager_outputs = eager_model(**inputs)
        self.assertTrue(blade_outputs)
        for x in blade_outputs:
            torch.testing.assert_close(
                blade_outputs[x].float(), eager_outputs[x].float(), rtol=0.001, atol=0.001)
//  [python]
assert reshaped_tensordict["b"].shape == torch.Size([12, 5])

##############################################################################
# Splitting a ``TensorDict``
# --------------------------
#
# :meth:`TensorDict.split <tensordict.TensorDict.split>` is similar to
# :meth:`torch.Tensor.split`. It splits the :class:`~.TensorDict` into chunks. Each
# chunk is a :class:`~.TensorDict` with the same structure as the original one, but
# whose entries are views of the corresponding entries in the original
# :class:`~.TensorDict`.

chunks = tensordict.split([3, 1], dim=1)
assert chunks[0].batch_size == torch.Size([3, 3])
assert chunks[1].batch_size == torch.Size([3, 1])
torch.testing.assert_allclose(chunks[0]["a"], tensordict["a"][:, :-1])
//  [python]
assert reshaped_tensordict["b"].shape == torch.Size([12, 5])

##############################################################################
# Splitting a ``TensorDict``
# --------------------------
#
# :meth:`TensorDict.split <tensordict.TensorDict.split>` is similar to
# :meth:`torch.Tensor.split`. It splits the :class:`~.TensorDict` into chunks. Each
# chunk is a :class:`~.TensorDict` with the same structure as the original one, but
# whose entries are views of the corresponding entries in the original
# :class:`~.TensorDict`.

chunks = tensordict.split([3, 1], dim=1)
assert chunks[0].batch_size == torch.Size([3, 3])
assert chunks[1].batch_size == torch.Size([3, 1])
torch.testing.assert_close(chunks[0]["a"], tensordict["a"][:, :-1])
//  [python]
def _compare_script_and_mobile(self,
                                   model: torch.nn.Module,
                                   input: torch.Tensor):
        # Compares the numerical outputs for script and lite modules
        qengine = "qnnpack"
        with override_quantized_engine(qengine):
            script_module = torch.jit.script(model)
            script_module_result = script_module(input)

            max_retry = 5
            for retry in range(1, max_retry + 1):
                # retries `max_retry` times; breaks iff succeeds else throws exception
                try:
                    buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
                    buffer.seek(0)
                    mobile_module = _load_for_lite_interpreter(buffer)

                    mobile_module_result = mobile_module(input)

                    t.assert_close(script_module_result, mobile_module_result)
                    mobile_module_forward_result = mobile_module.forward(input)
                    t.assert_close(script_module_result, mobile_module_forward_result)

                    mobile_module_run_method_result = mobile_module.run_method("forward", input)
                    t.assert_close(script_module_result, mobile_module_run_method_result)
                except AssertionError as e:
                    if retry == max_retry:
                        raise e
                    else:
                        continue
                break
//  [python]
def _compare_script_and_mobile(self,
                                   model: torch.nn.Module,
                                   input: torch.Tensor):
        # Compares the numerical outputs for script and lite modules
        qengine = "qnnpack"
        with override_quantized_engine(qengine):
            script_module = torch.jit.script(model)
            script_module_result = script_module(input)

            max_retry = 5
            for retry in range(1, max_retry + 1):
                # retries `max_retry` times; breaks iff succeeds else throws exception
                try:
                    buffer = io.BytesIO(script_module._save_to_buffer_for_lite_interpreter())
                    buffer.seek(0)
                    mobile_module = _load_for_lite_interpreter(buffer)

                    mobile_module_result = mobile_module(input)

                    torch.testing.assert_close(script_module_result, mobile_module_result)
                    mobile_module_forward_result = mobile_module.forward(input)
                    torch.testing.assert_close(script_module_result, mobile_module_forward_result)

                    mobile_module_run_method_result = mobile_module.run_method("forward", input)
                    torch.testing.assert_close(script_module_result, mobile_module_run_method_result)
                except AssertionError as e:
                    if retry == max_retry:
                        raise e
                    else:
                        continue
                break
//  [python]
def test_nms_ref(self, iou, seed):
        torch.random.manual_seed(seed)
        err_msg = "NMS incompatible between CPU and reference implementation for IoU={}"
        boxes, scores = self._create_tensors_with_iou(1000, iou)
        keep_ref = self._reference_nms(boxes, scores, iou)
        keep = ops.nms(boxes, scores, iou)
        torch.testing.assert_allclose(keep, keep_ref, msg=err_msg.format(iou))
//  [python]
def test_nms_ref(self, iou, seed):
        torch.random.manual_seed(seed)
        err_msg = "NMS incompatible between CPU and reference implementation for IoU={}"
        boxes, scores = self._create_tensors_with_iou(1000, iou)
        keep_ref = self._reference_nms(boxes, scores, iou)
        keep = ops.nms(boxes, scores, iou)
        torch.testing.assert_close(keep, keep_ref, msg=err_msg.format(iou))
//  [python]
def test_torch_fx_trace(self, device, x_dtype=torch.float, rois_dtype=torch.float):
        op_obj = self.make_obj().to(device=device)
        graph_module = torch.fx.symbolic_trace(op_obj)
        pool_size = 5
        n_channels = 2 * (pool_size**2)
        x = torch.rand(2, n_channels, 5, 5, dtype=x_dtype, device=device)
        rois = torch.tensor(
            [[0, 0, 0, 9, 9], [0, 0, 5, 4, 9], [0, 5, 5, 9, 9], [1, 0, 0, 9, 9]],  # format is (xyxy)
            dtype=rois_dtype,
            device=device,
        )
        output_gt = op_obj(x, rois)
        assert output_gt.dtype == x.dtype
        output_fx = graph_module(x, rois)
        assert output_fx.dtype == x.dtype
        tol = 1e-5
        torch.allclose(output_gt, output_fx, rtol=tol, atol=tol)
//  [python]
def test_torch_fx_trace(self, device, x_dtype=torch.float, rois_dtype=torch.float):
        op_obj = self.make_obj().to(device=device)
        graph_module = torch.fx.symbolic_trace(op_obj)
        pool_size = 5
        n_channels = 2 * (pool_size**2)
        x = torch.rand(2, n_channels, 5, 5, dtype=x_dtype, device=device)
        rois = torch.tensor(
            [[0, 0, 0, 9, 9], [0, 0, 5, 4, 9], [0, 5, 5, 9, 9], [1, 0, 0, 9, 9]],  # format is (xyxy)
            dtype=rois_dtype,
            device=device,
        )
        output_gt = op_obj(x, rois)
        assert output_gt.dtype == x.dtype
        output_fx = graph_module(x, rois)
        assert output_fx.dtype == x.dtype
        tol = 1e-5
        torch.testing.assert_close(output_gt, output_fx, rtol=tol, atol=tol)
//  [python]
def test_self_multihead_attn_norm_add(self):
        grads = torch.randn_like(self.tst_inputs)

        for _ in range(0, 5):
            ref_outputs, _ = self.ref_layer.forward(
                self.ref_inputs,
                self.ref_inputs,
                self.ref_inputs,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            tst_outputs, _ = self.tst_layer.forward(
                self.tst_inputs,
                self.tst_inputs,
                self.tst_inputs,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            self.ref_inputs.backward(grads)
            self.tst_inputs.backward(grads)

        torch.allclose(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_self_multihead_attn_norm_add(self):
        grads = torch.randn_like(self.tst_inputs)

        for _ in range(0, 5):
            ref_outputs, _ = self.ref_layer.forward(
                self.ref_inputs,
                self.ref_inputs,
                self.ref_inputs,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            tst_outputs, _ = self.tst_layer.forward(
                self.tst_inputs,
                self.tst_inputs,
                self.tst_inputs,
                key_padding_mask=None,
                need_weights=False,
                attn_mask=None,
                is_training=True,
            )

            self.ref_inputs.backward(grads)
            self.tst_inputs.backward(grads)

        torch.testing.assert_close(self.ref_inputs, self.tst_inputs, atol=1e-5, rtol=1e-5)
        torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
        torch.testing.assert_close(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3)
//  [python]
def test_autocast_cpu(self, iou, dtype):
        boxes, scores = self._create_tensors_with_iou(1000, iou)
        with torch.cpu.amp.autocast():
            keep_ref_float = ops.nms(boxes.to(dtype).float(), scores.to(dtype).float(), iou)
            keep_dtype = ops.nms(boxes.to(dtype), scores.to(dtype), iou)
        torch.testing.assert_allclose(keep_ref_float, keep_dtype)
//  [python]
def test_autocast_cpu(self, iou, dtype):
        boxes, scores = self._create_tensors_with_iou(1000, iou)
        with torch.cpu.amp.autocast():
            keep_ref_float = ops.nms(boxes.to(dtype).float(), scores.to(dtype).float(), iou)
            keep_dtype = ops.nms(boxes.to(dtype), scores.to(dtype), iou)
        torch.testing.assert_close(keep_ref_float, keep_dtype)
//  [python]
def test_qnms(self, iou, scale, zero_point):
        # Note: we compare qnms vs nms instead of qnms vs reference implementation.
        # This is because with the int conversion, the trick used in _create_tensors_with_iou
        # doesn't really work (in fact, nms vs reference implem will also fail with ints)
        err_msg = "NMS and QNMS give different results for IoU={}"
        boxes, scores = self._create_tensors_with_iou(1000, iou)
        scores *= 100  # otherwise most scores would be 0 or 1 after int conversion

        qboxes = torch.quantize_per_tensor(boxes, scale=scale, zero_point=zero_point, dtype=torch.quint8)
        qscores = torch.quantize_per_tensor(scores, scale=scale, zero_point=zero_point, dtype=torch.quint8)

        boxes = qboxes.dequantize()
        scores = qscores.dequantize()

        keep = ops.nms(boxes, scores, iou)
        qkeep = ops.nms(qboxes, qscores, iou)

        torch.testing.assert_allclose(qkeep, keep, msg=err_msg.format(iou))
//  [python]
def test_qnms(self, iou, scale, zero_point):
        # Note: we compare qnms vs nms instead of qnms vs reference implementation.
        # This is because with the int conversion, the trick used in _create_tensors_with_iou
        # doesn't really work (in fact, nms vs reference implem will also fail with ints)
        err_msg = "NMS and QNMS give different results for IoU={}"
        boxes, scores = self._create_tensors_with_iou(1000, iou)
        scores *= 100  # otherwise most scores would be 0 or 1 after int conversion

        qboxes = torch.quantize_per_tensor(boxes, scale=scale, zero_point=zero_point, dtype=torch.quint8)
        qscores = torch.quantize_per_tensor(scores, scale=scale, zero_point=zero_point, dtype=torch.quint8)

        boxes = qboxes.dequantize()
        scores = qscores.dequantize()

        keep = ops.nms(boxes, scores, iou)
        qkeep = ops.nms(qboxes, qscores, iou)

        torch.testing.assert_close(qkeep, keep, msg=err_msg.format(iou))
//  [python]
def area_check(self, box, expected, atol=1e-4):
        out = ops.box_area(box)
        t.assert_allclose(out, expected, rtol=0.0, check_dtype=False, atol=atol)
//  [python]
def area_check(self, box, expected, atol=1e-4):
        out = ops.box_area(box)
        torch.testing.assert_close(out, expected, rtol=0.0, check_dtype=False, atol=atol)
//  [python]
def _test_fsdp_not_all_outputs_used_in_loss(
        self, sharding_strategy: ShardingStrategy
    ):
        class MyModule(nn.Module):
            def __init__(self):
                super().__init__()
                self.lin1 = nn.Linear(4, 4)
                self.lin2 = nn.Linear(4, 4)

            def forward(self, x):
                a = self.lin1(x)
                b = self.lin2(x)
                return (a, b)

        def _check_resharded(fsdp_module):
            handle = fsdp_module._handle
            if not handle:
                return
            param = handle.flat_param
            if handle.uses_sharded_strategy:
                full_param = param._full_param_padded
                self.assertEqual(full_param.storage().size(), 0)

            self.assertEqual(param.data_ptr(), param._local_shard.data_ptr())

        def _check_equal(local, fsdp):
            with FSDP.summon_full_params(fsdp):
                for p1, p2 in zip(fsdp.parameters(), local.parameters()):
                    torch.testing.assert_allclose(p1, p2)

        fsdp_ctor = functools.partial(FSDP, sharding_strategy=sharding_strategy)
        m = MyModule().cuda()
        m_local = deepcopy(m)
        local_m = m_local
        prev_params = [p.clone() for p in m_local.parameters()]

        m.lin1 = fsdp_ctor(m.lin1)
        m = fsdp_ctor(m)
        _check_equal(m_local, m)

        opt = torch.optim.SGD(m.parameters(), lr=1e-3)
        opt_local = torch.optim.SGD(local_m.parameters(), lr=1e-3)

        for i in range(6):
            t = torch.ones(4, device="cuda")
            a, b = m(t)
            local_a, local_b = local_m(t)
            if i < 2:
                # use both params in loss computation. Later,
                # b will go unused and we check grads are the
                # same as local training.
                loss = (a @ b).sum()
                loss_local = (local_a @ local_b).sum()
            else:
                loss = a.sum()
                loss_local = local_a.sum()

            loss.backward()
            loss_local.backward()
            _check_resharded(m)
            opt.step()
            opt_local.step()
            _check_equal(m_local, m)
            # Ensure at least some change from previous params, otherwise
            # above check would be vacuously true.
            self.assertTrue(
                any(
                    not torch.equal(p1, p2)
                    for p1, p2 in zip(prev_params, m_local.parameters())
                )
            )
            prev_params = [p.clone() for p in local_m.parameters()]
            opt.zero_grad()
            opt_local.zero_grad()

        dist.barrier()
//  [python]
def _test_fsdp_not_all_outputs_used_in_loss(
        self, sharding_strategy: ShardingStrategy
    ):
        class MyModule(nn.Module):
            def __init__(self):
                super().__init__()
                self.lin1 = nn.Linear(4, 4)
                self.lin2 = nn.Linear(4, 4)

            def forward(self, x):
                a = self.lin1(x)
                b = self.lin2(x)
                return (a, b)

        def _check_resharded(fsdp_module):
            handle = fsdp_module._handle
            if not handle:
                return
            param = handle.flat_param
            if handle.uses_sharded_strategy:
                full_param = param._full_param_padded
                self.assertEqual(full_param.storage().size(), 0)

            self.assertEqual(param.data_ptr(), param._local_shard.data_ptr())

        def _check_equal(local, fsdp):
            with FSDP.summon_full_params(fsdp):
                for p1, p2 in zip(fsdp.parameters(), local.parameters()):
                    torch.testing.assert_close(p1, p2)

        fsdp_ctor = functools.partial(FSDP, sharding_strategy=sharding_strategy)
        m = MyModule().cuda()
        m_local = deepcopy(m)
        local_m = m_local
        prev_params = [p.clone() for p in m_local.parameters()]

        m.lin1 = fsdp_ctor(m.lin1)
        m = fsdp_ctor(m)
        _check_equal(m_local, m)

        opt = torch.optim.SGD(m.parameters(), lr=1e-3)
        opt_local = torch.optim.SGD(local_m.parameters(), lr=1e-3)

        for i in range(6):
            t = torch.ones(4, device="cuda")
            a, b = m(t)
            local_a, local_b = local_m(t)
            if i < 2:
                # use both params in loss computation. Later,
                # b will go unused and we check grads are the
                # same as local training.
                loss = (a @ b).sum()
                loss_local = (local_a @ local_b).sum()
            else:
                loss = a.sum()
                loss_local = local_a.sum()

            loss.backward()
            loss_local.backward()
            _check_resharded(m)
            opt.step()
            opt_local.step()
            _check_equal(m_local, m)
            # Ensure at least some change from previous params, otherwise
            # above check would be vacuously true.
            self.assertTrue(
                any(
                    not torch.equal(p1, p2)
                    for p1, p2 in zip(prev_params, m_local.parameters())
                )
            )
            prev_params = [p.clone() for p in local_m.parameters()]
            opt.zero_grad()
            opt_local.zero_grad()

        dist.barrier()
//  [python]
def test_batched_nms_implementations(self, seed):
        """Make sure that both implementations of batched_nms yield identical results"""
        torch.random.manual_seed(seed)

        num_boxes = 1000
        iou_threshold = 0.9

        boxes = torch.cat((torch.rand(num_boxes, 2), torch.rand(num_boxes, 2) + 10), dim=1)
        assert max(boxes[:, 0]) < min(boxes[:, 2])  # x1 < x2
        assert max(boxes[:, 1]) < min(boxes[:, 3])  # y1 < y2

        scores = torch.rand(num_boxes)
        idxs = torch.randint(0, 4, size=(num_boxes,))
        keep_vanilla = ops.boxes._batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
        keep_trick = ops.boxes._batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)

        torch.testing.assert_allclose(
            keep_vanilla, keep_trick, msg="The vanilla and the trick implementation yield different nms outputs."
        )

        # Also make sure an empty tensor is returned if boxes is empty
        empty = torch.empty((0,), dtype=torch.int64)
        torch.testing.assert_allclose(empty, ops.batched_nms(empty, None, None, None))
//  [python]
def test_batched_nms_implementations(self, seed):
        """Make sure that both implementations of batched_nms yield identical results"""
        torch.random.manual_seed(seed)

        num_boxes = 1000
        iou_threshold = 0.9

        boxes = torch.cat((torch.rand(num_boxes, 2), torch.rand(num_boxes, 2) + 10), dim=1)
        assert max(boxes[:, 0]) < min(boxes[:, 2])  # x1 < x2
        assert max(boxes[:, 1]) < min(boxes[:, 3])  # y1 < y2

        scores = torch.rand(num_boxes)
        idxs = torch.randint(0, 4, size=(num_boxes,))
        keep_vanilla = ops.boxes._batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
        keep_trick = ops.boxes._batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)

        torch.testing.assert_close(
            keep_vanilla, keep_trick, msg="The vanilla and the trick implementation yield different nms outputs."
        )

        # Also make sure an empty tensor is returned if boxes is empty
        empty = torch.empty((0,), dtype=torch.int64)
        torch.testing.assert_close(empty, ops.batched_nms(empty, None, None, None))
//  [python]
def test_box_area_jit(self):
        box_tensor = torch.tensor([[0, 0, 100, 100], [0, 0, 0, 0]], dtype=torch.float)
        expected = ops.box_area(box_tensor)
        scripted_fn = torch.jit.script(ops.box_area)
        scripted_area = scripted_fn(box_tensor)
        torch.testing.assert_allclose(scripted_area, expected)
//  [python]
def test_box_area_jit(self):
        box_tensor = torch.tensor([[0, 0, 100, 100], [0, 0, 0, 0]], dtype=torch.float)
        expected = ops.box_area(box_tensor)
        scripted_fn = torch.jit.script(ops.box_area)
        scripted_area = scripted_fn(box_tensor)
        torch.testing.assert_close(scripted_area, expected)
//  [python]
def test_bbox_convert_jit(self):
        box_tensor = torch.tensor(
            [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float
        )

        scripted_fn = torch.jit.script(ops.box_convert)

        box_xywh = ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="xywh")
        scripted_xywh = scripted_fn(box_tensor, "xyxy", "xywh")
        torch.testing.assert_allclose(scripted_xywh, box_xywh)

        box_cxcywh = ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="cxcywh")
        scripted_cxcywh = scripted_fn(box_tensor, "xyxy", "cxcywh")
        torch.testing.assert_allclose(scripted_cxcywh, box_cxcywh)
//  [python]
def test_bbox_convert_jit(self):
        box_tensor = torch.tensor(
            [[0, 0, 100, 100], [0, 0, 0, 0], [10, 15, 30, 35], [23, 35, 93, 95]], dtype=torch.float
        )

        scripted_fn = torch.jit.script(ops.box_convert)

        box_xywh = ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="xywh")
        scripted_xywh = scripted_fn(box_tensor, "xyxy", "xywh")
        torch.testing.assert_close(scripted_xywh, box_xywh)

        box_cxcywh = ops.box_convert(box_tensor, in_fmt="xyxy", out_fmt="cxcywh")
        scripted_cxcywh = scripted_fn(box_tensor, "xyxy", "cxcywh")
        torch.testing.assert_close(scripted_cxcywh, box_cxcywh)
//  [python]
def test_forward(self, device, contiguous, batch_sz, dtype=None):
        dtype = dtype or self.dtype
        x, _, offset, mask, _, stride, padding, dilation = self.get_fn_args(device, contiguous, batch_sz, dtype)
        in_channels = 6
        out_channels = 2
        kernel_size = (3, 2)
        groups = 2
        tol = 2e-3 if dtype is torch.half else 1e-5

        layer = self.make_obj(in_channels, out_channels, kernel_size, groups, wrap=False).to(
            device=x.device, dtype=dtype
        )
        res = layer(x, offset, mask)

        weight = layer.weight.data
        bias = layer.bias.data
        expected = self.expected_fn(x, weight, offset, mask, bias, stride=stride, padding=padding, dilation=dilation)

        torch.testing.assert_allclose(
            res.to(expected), expected, rtol=tol, atol=tol, msg=f"\nres:\n{res}\nexpected:\n{expected}"
        )

        # no modulation test
        res = layer(x, offset)
        expected = self.expected_fn(x, weight, offset, None, bias, stride=stride, padding=padding, dilation=dilation)

        torch.testing.assert_allclose(
            res.to(expected), expected, rtol=tol, atol=tol, msg=f"\nres:\n{res}\nexpected:\n{expected}"
        )
//  [python]
def test_forward(self, device, contiguous, batch_sz, dtype=None):
        dtype = dtype or self.dtype
        x, _, offset, mask, _, stride, padding, dilation = self.get_fn_args(device, contiguous, batch_sz, dtype)
        in_channels = 6
        out_channels = 2
        kernel_size = (3, 2)
        groups = 2
        tol = 2e-3 if dtype is torch.half else 1e-5

        layer = self.make_obj(in_channels, out_channels, kernel_size, groups, wrap=False).to(
            device=x.device, dtype=dtype
        )
        res = layer(x, offset, mask)

        weight = layer.weight.data
        bias = layer.bias.data
        expected = self.expected_fn(x, weight, offset, mask, bias, stride=stride, padding=padding, dilation=dilation)

        torch.testing.assert_close(
            res.to(expected), expected, rtol=tol, atol=tol, msg=f"\nres:\n{res}\nexpected:\n{expected}"
        )

        # no modulation test
        res = layer(x, offset)
        expected = self.expected_fn(x, weight, offset, None, bias, stride=stride, padding=padding, dilation=dilation)

        torch.testing.assert_close(
            res.to(expected), expected, rtol=tol, atol=tol, msg=f"\nres:\n{res}\nexpected:\n{expected}"
        )
//  [python]
def _run_test(target_fn: Callable, actual_box1, actual_box2, dtypes, atol, expected):
        for dtype in dtypes:
            actual_box1 = torch.tensor(actual_box1, dtype=dtype)
            actual_box2 = torch.tensor(actual_box2, dtype=dtype)
            expected_box = torch.tensor(expected)
            out = target_fn(actual_box1, actual_box2)
            torch.testing.assert_allclose(out, expected_box, rtol=0.0, check_dtype=False, atol=atol)
//  [python]
def _run_test(target_fn: Callable, actual_box1, actual_box2, dtypes, atol, expected):
        for dtype in dtypes:
            actual_box1 = torch.tensor(actual_box1, dtype=dtype)
            actual_box2 = torch.tensor(actual_box2, dtype=dtype)
            expected_box = torch.tensor(expected)
            out = target_fn(actual_box1, actual_box2)
            torch.testing.assert_close(out, expected_box, rtol=0.0, check_dtype=False, atol=atol)
//  [python]
def test_frozenbatchnorm2d_eps(self, seed):
        torch.random.manual_seed(seed)
        sample_size = (4, 32, 28, 28)
        x = torch.rand(sample_size)
        state_dict = dict(
            weight=torch.rand(sample_size[1]),
            bias=torch.rand(sample_size[1]),
            running_mean=torch.rand(sample_size[1]),
            running_var=torch.rand(sample_size[1]),
            num_batches_tracked=torch.tensor(100),
        )

        # Check that default eps is equal to the one of BN
        fbn = ops.misc.FrozenBatchNorm2d(sample_size[1])
        fbn.load_state_dict(state_dict, strict=False)
        bn = torch.nn.BatchNorm2d(sample_size[1]).eval()
        bn.load_state_dict(state_dict)
        # Difference is expected to fall in an acceptable range
        torch.testing.assert_allclose(fbn(x), bn(x), rtol=1e-5, atol=1e-6)

        # Check computation for eps > 0
        fbn = ops.misc.FrozenBatchNorm2d(sample_size[1], eps=1e-5)
        fbn.load_state_dict(state_dict, strict=False)
        bn = torch.nn.BatchNorm2d(sample_size[1], eps=1e-5).eval()
        bn.load_state_dict(state_dict)
        torch.testing.assert_allclose(fbn(x), bn(x), rtol=1e-5, atol=1e-6)
//  [python]
def test_frozenbatchnorm2d_eps(self, seed):
        torch.random.manual_seed(seed)
        sample_size = (4, 32, 28, 28)
        x = torch.rand(sample_size)
        state_dict = dict(
            weight=torch.rand(sample_size[1]),
            bias=torch.rand(sample_size[1]),
            running_mean=torch.rand(sample_size[1]),
            running_var=torch.rand(sample_size[1]),
            num_batches_tracked=torch.tensor(100),
        )

        # Check that default eps is equal to the one of BN
        fbn = ops.misc.FrozenBatchNorm2d(sample_size[1])
        fbn.load_state_dict(state_dict, strict=False)
        bn = torch.nn.BatchNorm2d(sample_size[1]).eval()
        bn.load_state_dict(state_dict)
        # Difference is expected to fall in an acceptable range
        torch.testing.assert_close(fbn(x), bn(x), rtol=1e-5, atol=1e-6)

        # Check computation for eps > 0
        fbn = ops.misc.FrozenBatchNorm2d(sample_size[1], eps=1e-5)
        fbn.load_state_dict(state_dict, strict=False)
        bn = torch.nn.BatchNorm2d(sample_size[1], eps=1e-5).eval()
        bn.load_state_dict(state_dict)
        torch.testing.assert_close(fbn(x), bn(x), rtol=1e-5, atol=1e-6)
//  [python]
def _run_jit_test(target_fn: Callable, actual_box: List):
        box_tensor = torch.tensor(actual_box, dtype=torch.float)
        expected = target_fn(box_tensor, box_tensor)
        scripted_fn = torch.jit.script(target_fn)
        scripted_out = scripted_fn(box_tensor, box_tensor)
        torch.testing.assert_allclose(scripted_out, expected)
//  [python]
def _run_jit_test(target_fn: Callable, actual_box: List):
        box_tensor = torch.tensor(actual_box, dtype=torch.float)
        expected = target_fn(box_tensor, box_tensor)
        scripted_fn = torch.jit.script(target_fn)
        scripted_out = scripted_fn(box_tensor, box_tensor)
        torch.testing.assert_close(scripted_out, expected)
//  [python]
def _run_cartesian_test(target_fn: Callable):
        boxes1 = gen_box(5)
        boxes2 = gen_box(7)
        a = TestIouBase._cartesian_product(boxes1, boxes2, target_fn)
        b = target_fn(boxes1, boxes2)
        t.assert_close(a, b)
//  [python]
def _run_cartesian_test(target_fn: Callable):
        boxes1 = gen_box(5)
        boxes2 = gen_box(7)
        a = TestIouBase._cartesian_product(boxes1, boxes2, target_fn)
        b = target_fn(boxes1, boxes2)
        torch.testing.assert_close(a, b)
//  [python]
def assert_iou_loss(iou_fn, box1, box2, expected_loss, device, reduction="none"):
    computed_loss = iou_fn(box1, box2, reduction=reduction)
    expected_loss = torch.tensor(expected_loss, device=device)
    torch.testing.assert_allclose(computed_loss, expected_loss)
//  [python]
def assert_iou_loss(iou_fn, box1, box2, expected_loss, device, reduction="none"):
    computed_loss = iou_fn(box1, box2, reduction=reduction)
    expected_loss = torch.tensor(expected_loss, device=device)
    torch.testing.assert_close(computed_loss, expected_loss)
//  [python]
def assert_empty_loss(iou_fn, dtype, device):
    box1 = torch.randn([0, 4], dtype=dtype, device=device).requires_grad_()
    box2 = torch.randn([0, 4], dtype=dtype, device=device).requires_grad_()
    loss = iou_fn(box1, box2, reduction="mean")
    loss.backward()
    torch.testing.assert_allclose(loss, torch.tensor(0.0, device=device))
    assert box1.grad is not None, "box1.grad should not be None after backward is called"
    assert box2.grad is not None, "box2.grad should not be None after backward is called"
    loss = iou_fn(box1, box2, reduction="none")
    assert loss.numel() == 0, f"{str(iou_fn)} for two empty box should be empty"
//  [python]
def assert_empty_loss(iou_fn, dtype, device):
    box1 = torch.randn([0, 4], dtype=dtype, device=device).requires_grad_()
    box2 = torch.randn([0, 4], dtype=dtype, device=device).requires_grad_()
    loss = iou_fn(box1, box2, reduction="mean")
    loss.backward()
    torch.testing.assert_close(loss, torch.tensor(0.0, device=device))
    assert box1.grad is not None, "box1.grad should not be None after backward is called"
    assert box2.grad is not None, "box2.grad should not be None after backward is called"
    loss = iou_fn(box1, box2, reduction="none")
    assert loss.numel() == 0, f"{str(iou_fn)} for two empty box should be empty"
//  [python]
def test_jit(self, alpha, gamma, reduction, device, dtype, seed):
        if device == "cpu" and dtype is torch.half:
            pytest.skip("Currently torch.half is not fully supported on cpu")
        script_fn = torch.jit.script(ops.sigmoid_focal_loss)
        torch.random.manual_seed(seed)
        inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device)
        focal_loss = ops.sigmoid_focal_loss(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction)
        scripted_focal_loss = script_fn(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction)

        tol = 1e-3 if dtype is torch.half else 1e-5
        torch.testing.assert_allclose(focal_loss, scripted_focal_loss, rtol=tol, atol=tol)
//  [python]
def test_jit(self, alpha, gamma, reduction, device, dtype, seed):
        if device == "cpu" and dtype is torch.half:
            pytest.skip("Currently torch.half is not fully supported on cpu")
        script_fn = torch.jit.script(ops.sigmoid_focal_loss)
        torch.random.manual_seed(seed)
        inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device)
        focal_loss = ops.sigmoid_focal_loss(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction)
        scripted_focal_loss = script_fn(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction)

        tol = 1e-3 if dtype is torch.half else 1e-5
        torch.testing.assert_close(focal_loss, scripted_focal_loss, rtol=tol, atol=tol)
//  [python]
def test_equal_ce_loss(self, reduction, device, dtype, seed):
        if device == "cpu" and dtype is torch.half:
            pytest.skip("Currently torch.half is not fully supported on cpu")
        # focal loss should be equal ce_loss if alpha=-1 and gamma=0
        alpha = -1
        gamma = 0
        torch.random.manual_seed(seed)
        inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device)
        inputs_fl = inputs.clone().requires_grad_()
        targets_fl = targets.clone()
        inputs_ce = inputs.clone().requires_grad_()
        targets_ce = targets.clone()
        focal_loss = ops.sigmoid_focal_loss(inputs_fl, targets_fl, gamma=gamma, alpha=alpha, reduction=reduction)
        ce_loss = F.binary_cross_entropy_with_logits(inputs_ce, targets_ce, reduction=reduction)

        torch.testing.assert_allclose(focal_loss, ce_loss)

        focal_loss.backward()
        ce_loss.backward()
        torch.testing.assert_allclose(inputs_fl.grad, inputs_ce.grad)
//  [python]
def test_equal_ce_loss(self, reduction, device, dtype, seed):
        if device == "cpu" and dtype is torch.half:
            pytest.skip("Currently torch.half is not fully supported on cpu")
        # focal loss should be equal ce_loss if alpha=-1 and gamma=0
        alpha = -1
        gamma = 0
        torch.random.manual_seed(seed)
        inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device)
        inputs_fl = inputs.clone().requires_grad_()
        targets_fl = targets.clone()
        inputs_ce = inputs.clone().requires_grad_()
        targets_ce = targets.clone()
        focal_loss = ops.sigmoid_focal_loss(inputs_fl, targets_fl, gamma=gamma, alpha=alpha, reduction=reduction)
        ce_loss = F.binary_cross_entropy_with_logits(inputs_ce, targets_ce, reduction=reduction)

        torch.testing.assert_close(focal_loss, ce_loss)

        focal_loss.backward()
        ce_loss.backward()
        torch.testing.assert_close(inputs_fl.grad, inputs_ce.grad)
//  [python]
def test_correct_ratio(self, alpha, gamma, device, dtype, seed):
        if device == "cpu" and dtype is torch.half:
            pytest.skip("Currently torch.half is not fully supported on cpu")
        # For testing the ratio with manual calculation, we require the reduction to be "none"
        reduction = "none"
        torch.random.manual_seed(seed)
        inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device)
        focal_loss = ops.sigmoid_focal_loss(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction)
        ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction=reduction)

        assert torch.all(
            focal_loss <= ce_loss
        ), "focal loss must be less or equal to cross entropy loss with same input"

        loss_ratio = (focal_loss / ce_loss).squeeze()
        prob = torch.sigmoid(inputs)
        p_t = prob * targets + (1 - prob) * (1 - targets)
        correct_ratio = (1.0 - p_t) ** gamma
        if alpha >= 0:
            alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
            correct_ratio = correct_ratio * alpha_t

        tol = 1e-3 if dtype is torch.half else 1e-5
        torch.testing.assert_allclose(correct_ratio, loss_ratio, atol=tol, rtol=tol)
//  [python]
def test_correct_ratio(self, alpha, gamma, device, dtype, seed):
        if device == "cpu" and dtype is torch.half:
            pytest.skip("Currently torch.half is not fully supported on cpu")
        # For testing the ratio with manual calculation, we require the reduction to be "none"
        reduction = "none"
        torch.random.manual_seed(seed)
        inputs, targets = self._generate_diverse_input_target_pair(dtype=dtype, device=device)
        focal_loss = ops.sigmoid_focal_loss(inputs, targets, gamma=gamma, alpha=alpha, reduction=reduction)
        ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction=reduction)

        assert torch.all(
            focal_loss <= ce_loss
        ), "focal loss must be less or equal to cross entropy loss with same input"

        loss_ratio = (focal_loss / ce_loss).squeeze()
        prob = torch.sigmoid(inputs)
        p_t = prob * targets + (1 - prob) * (1 - targets)
        correct_ratio = (1.0 - p_t) ** gamma
        if alpha >= 0:
            alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
            correct_ratio = correct_ratio * alpha_t

        tol = 1e-3 if dtype is torch.half else 1e-5
        torch.testing.assert_close(correct_ratio, loss_ratio, atol=tol, rtol=tol)
//  [python]
def test_compare_fused_obs_fq_oss_module(self, device):
        mod = FusedMovingAvgObsFakeQuantize()
        torch.ao.quantization.enable_fake_quant(mod)
        torch.ao.quantization.enable_observer(mod)
        mod.to(device)

        mod_ref = FakeQuantize()
        torch.ao.quantization.enable_fake_quant(mod_ref)
        torch.ao.quantization.enable_observer(mod_ref)
        mod_ref.to(device)

        for i in range(10):
            x = torch.randn(5, 5, device=device)
            out = mod(x)
            out_ref = mod_ref(x)
            self.assertEqual(out, out_ref)
            self.assertEqual(
                mod_ref.activation_post_process.min_val,
                mod.activation_post_process.min_val,
            )
            self.assertEqual(
                mod_ref.activation_post_process.max_val,
                mod.activation_post_process.max_val,
            )
//  [python]
def test_compare_fused_obs_fq_oss_module(self, device):
        mod = FusedMovingAvgObsFakeQuantize()
        torch.ao.quantization.enable_fake_quant(mod)
        torch.ao.quantization.enable_observer(mod)
        mod.to(device)

        mod_ref = FakeQuantize()
        torch.ao.quantization.enable_fake_quant(mod_ref)
        torch.ao.quantization.enable_observer(mod_ref)
        mod_ref.to(device)

        for i in range(10):
            x = torch.randn(5, 5, device=device)
            out = mod(x)
            out_ref = mod_ref(x)
            torch.testing.assert_close(out, out_ref)
            torch.testing.assert_close(
                mod_ref.activation_post_process.min_val,
                mod.activation_post_process.min_val,
            )
            torch.testing.assert_close(
                mod_ref.activation_post_process.max_val,
                mod.activation_post_process.max_val,
            )