jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _test5d(self, align_corners):
        T, H, W = 3, 4, 5
        # Construct a grid to obtain indentity sampling
        input = torch.randn(H * W).view(1, 1, H, W).float()
        input = torch.stack([input, input + 1, input + 2], dim=2)
        coords = torch.meshgrid(torch.arange(T), torch.arange(W), torch.arange(H))
        coords = torch.stack(coords, dim=-1).float().permute(0, 2, 1, 3)[None]

        if not align_corners:
            coords = coords + 0.5
        sampled_input = bilinear_sampler(input, coords, align_corners=align_corners)
        torch.testing.assert_close(input, sampled_input)
//  [python]
def test_max_autotune_cutlass_backend_addmm(
        self, dynamic, max_autotune_gemm_backends
    ):
        """
        Make sure autotuning addmm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def addmm(x, a, b, alpha, beta):
            return torch.addmm(x, a, b, alpha=alpha, beta=beta)

        def compare_results(
            m: int, k: int, n: int, alpha: float, beta: float, x_shape: List[int]
        ) -> None:
            x = torch.randn(x_shape).cuda().half()
            a = torch.randn(m, k).cuda().half()
            b = torch.randn(k, n).cuda().half()
            y_expected = addmm(x, a, b, alpha, beta)

            compiled_fn = torch.compile(addmm, dynamic=dynamic)
            y = compiled_fn(x, a, b, alpha, beta)
            t.assert_close(y, y_expected)

        with config.patch(
            {
                "max_autotune": True,
                # Some Cutlass Kernels fail with IMA on this example, which leads to unrecoverable CUDA errors
                # unless we tune in a subproc here.
                "autotune_in_subproc": True,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 4,
            }
        ):
            # No broadcast
            compare_results(4096, 25728, 2048, 2.0, 0.4, [4096, 2048])
            # Broadcast first dim.
            compare_results(4096, 25728, 2048, 2.0, 0.4, [2048])
            # Broadcast last dim.
            compare_results(4096, 25728, 2048, 2.0, 0.4, [4096, 1])
//  [python]
def test_max_autotune_cutlass_backend_addmm(
        self, dynamic, max_autotune_gemm_backends
    ):
        """
        Make sure autotuning addmm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def addmm(x, a, b, alpha, beta):
            return torch.addmm(x, a, b, alpha=alpha, beta=beta)

        def compare_results(
            m: int, k: int, n: int, alpha: float, beta: float, x_shape: List[int]
        ) -> None:
            x = torch.randn(x_shape).cuda().half()
            a = torch.randn(m, k).cuda().half()
            b = torch.randn(k, n).cuda().half()
            y_expected = addmm(x, a, b, alpha, beta)

            compiled_fn = torch.compile(addmm, dynamic=dynamic)
            y = compiled_fn(x, a, b, alpha, beta)
            torch.testing.assert_close(y, y_expected)

        with config.patch(
            {
                "max_autotune": True,
                # Some Cutlass Kernels fail with IMA on this example, which leads to unrecoverable CUDA errors
                # unless we tune in a subproc here.
                "autotune_in_subproc": True,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 4,
            }
        ):
            # No broadcast
            compare_results(4096, 25728, 2048, 2.0, 0.4, [4096, 2048])
            # Broadcast first dim.
            compare_results(4096, 25728, 2048, 2.0, 0.4, [2048])
            # Broadcast last dim.
            compare_results(4096, 25728, 2048, 2.0, 0.4, [4096, 1])
//  [python]
def _test_binary_confusion_matrix_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        normalize: Optional[str] = None,
    ) -> None:
        sklearn_result = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize=normalize)
        ).to(torch.float32)

        torch.testing.assert_allclose(
            bcm(input, target, normalize=normalize).to(torch.float32),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_binary_confusion_matrix_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        normalize: Optional[str] = None,
    ) -> None:
        sklearn_result = torch.tensor(
            skcm(target, input, labels=[0, 1], normalize=normalize)
        ).to(torch.float32)

        torch.testing.assert_close(
            bcm(input, target, normalize=normalize).to(torch.float32),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def copy():
    session = brt.Session(alloc_func=caching_allocator_alloc, free_func=caching_allocator_delete)
    model_path = os.path.join(os.path.dirname(__file__), "arg_alias.mlir")
    session.load(model_path)
    req = session.new_request_context(torch.cuda.current_stream()._as_parameter_.value)

    inputs = []
    outputs = []

    for offset in session.get_input_arg_offsets():
        data = np.random.random(size=session.get_static_shape(offset))
        inputs.append(torch.tensor(data, dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, inputs[-1].data_ptr())

    for offset in session.get_output_arg_offsets():
        outputs.append(torch.empty(session.get_static_shape(offset), dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, outputs[-1].data_ptr())

    req.finish_io_binding()
    req.run()
    req.sync()

    t.assert_close(inputs[0], outputs[0])
    t.assert_close(inputs[1], outputs[1])
//  [python]
def copy():
    session = brt.Session(alloc_func=caching_allocator_alloc, free_func=caching_allocator_delete)
    model_path = os.path.join(os.path.dirname(__file__), "arg_alias.mlir")
    session.load(model_path)
    req = session.new_request_context(torch.cuda.current_stream()._as_parameter_.value)

    inputs = []
    outputs = []

    for offset in session.get_input_arg_offsets():
        data = np.random.random(size=session.get_static_shape(offset))
        inputs.append(torch.tensor(data, dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, inputs[-1].data_ptr())

    for offset in session.get_output_arg_offsets():
        outputs.append(torch.empty(session.get_static_shape(offset), dtype=torch.float32, device="cuda"))
        req.bind_arg(offset, outputs[-1].data_ptr())

    req.finish_io_binding()
    req.run()
    req.sync()

    torch.testing.assert_close(inputs[0], outputs[0])
    torch.testing.assert_close(inputs[1], outputs[1])
//  [python]
def test_click_through_rate_with_valid_input(self) -> None:
        input = torch.tensor([0, 1, 0, 1, 1, 0, 0, 1])
        weights = torch.tensor([1.0, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 2.0])
        t.assert_close(click_through_rate(input), torch.tensor(0.5))
        t.assert_close(
            click_through_rate(input, weights), torch.tensor(0.58333334)
        )

        input = torch.tensor([[0, 1, 0, 1], [1, 0, 0, 1]])
        weights = torch.tensor([[1.0, 2.0, 1.0, 2.0], [1.0, 2.0, 1.0, 1.0]])
        t.assert_close(
            click_through_rate(input, num_tasks=2), torch.tensor([0.5, 0.5])
        )
        t.assert_close(
            click_through_rate(input, weights, num_tasks=2),
            torch.tensor([0.66666667, 0.4]),
        )
//  [python]
def test_click_through_rate_with_valid_input(self) -> None:
        input = torch.tensor([0, 1, 0, 1, 1, 0, 0, 1])
        weights = torch.tensor([1.0, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 2.0])
        torch.testing.assert_close(click_through_rate(input), torch.tensor(0.5))
        torch.testing.assert_close(
            click_through_rate(input, weights), torch.tensor(0.58333334)
        )

        input = torch.tensor([[0, 1, 0, 1], [1, 0, 0, 1]])
        weights = torch.tensor([[1.0, 2.0, 1.0, 2.0], [1.0, 2.0, 1.0, 1.0]])
        torch.testing.assert_close(
            click_through_rate(input, num_tasks=2), torch.tensor([0.5, 0.5])
        )
        torch.testing.assert_close(
            click_through_rate(input, weights, num_tasks=2),
            torch.tensor([0.66666667, 0.4]),
        )
//  [python]
def _test_multiclass_confusion_matrix_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        normalize: Optional[str] = None,
    ) -> None:
        sklearn_result = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize=normalize)
        ).to(torch.float32)

        torch.testing.assert_allclose(
            mccm(input, target, num_classes=num_classes, normalize=normalize).to(
                torch.float32
            ),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_multiclass_confusion_matrix_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        normalize: Optional[str] = None,
    ) -> None:
        sklearn_result = torch.tensor(
            skcm(target, input, labels=list(range(num_classes)), normalize=normalize)
        ).to(torch.float32)

        torch.testing.assert_close(
            mccm(input, target, num_classes=num_classes, normalize=normalize).to(
                torch.float32
            ),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_binary_confusion_matrix_score_thresholding(self) -> None:
        num_classes = 2
        threshold = 0.7
        input = torch.tensor([0.7, 0.6, 0.5, 0.3, 0.9, 0.1, 1.0, 0.95, 0.2])
        input_thresholded = torch.tensor([1, 0, 0, 0, 1, 0, 1, 1, 0])
        target = torch.randint(high=num_classes, size=(9,))

        sklearn_result = torch.tensor(
            skcm(target, input_thresholded, labels=[0, 1])
        ).to(torch.float32)

        my_result = bcm(input, target, threshold=threshold).to(torch.float32)

        # test threshold larger than every prediction gives 0
        torch.allclose(
            my_result,
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_binary_confusion_matrix_score_thresholding(self) -> None:
        num_classes = 2
        threshold = 0.7
        input = torch.tensor([0.7, 0.6, 0.5, 0.3, 0.9, 0.1, 1.0, 0.95, 0.2])
        input_thresholded = torch.tensor([1, 0, 0, 0, 1, 0, 1, 1, 0])
        target = torch.randint(high=num_classes, size=(9,))

        sklearn_result = torch.tensor(
            skcm(target, input_thresholded, labels=[0, 1])
        ).to(torch.float32)

        my_result = bcm(input, target, threshold=threshold).to(torch.float32)

        # test threshold larger than every prediction gives 0
        torch.testing.assert_close(
            my_result,
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _eval_models(
        self,
        m1: DistributedModelParallel,
        m2: DistributedModelParallel,
        batch: ModelInput,
        is_deterministic: bool = True,
    ) -> None:
        with torch.no_grad():
            loss1, pred1 = m1(batch)
            loss2, pred2 = m2(batch)

        if is_deterministic:
            self.assertTrue(torch.equal(loss1, loss2))
            self.assertTrue(torch.equal(pred1, pred2))
        else:
            rtol, atol = _get_default_rtol_and_atol(loss1, loss2)
            self.assertEqual(loss1, loss2, rtol=rtol, atol=atol)
            rtol, atol = _get_default_rtol_and_atol(pred1, pred2)
            self.assertEqual(pred1, pred2, rtol=rtol, atol=atol)
//  [python]
def _eval_models(
        self,
        m1: DistributedModelParallel,
        m2: DistributedModelParallel,
        batch: ModelInput,
        is_deterministic: bool = True,
    ) -> None:
        with torch.no_grad():
            loss1, pred1 = m1(batch)
            loss2, pred2 = m2(batch)

        if is_deterministic:
            self.assertTrue(torch.equal(loss1, loss2))
            self.assertTrue(torch.equal(pred1, pred2))
        else:
            rtol, atol = _get_default_rtol_and_atol(loss1, loss2)
            torch.testing.assert_close(loss1, loss2, rtol=rtol, atol=atol)
            rtol, atol = _get_default_rtol_and_atol(pred1, pred2)
            torch.testing.assert_close(pred1, pred2, rtol=rtol, atol=atol)
//  [python]
def test_multiclass_confusion_matrix_with_probabilities(self) -> None:
        # should reduce to [2, 0, 1, 2, 1]
        input = torch.tensor(
            [
                [0.2948, 0.3343, 0.3709],
                [0.4988, 0.4836, 0.0176],
                [0.3727, 0.5145, 0.1128],
                [0.3759, 0.2115, 0.4126],
                [0.3076, 0.4226, 0.2698],
            ]
        )

        target = torch.tensor([0, 1, 2, 0, 0])

        expected_output = torch.tensor([[0, 1, 2], [1, 0, 0], [0, 1, 0]]).to(
            torch.float32
        )

        output = mccm(input, target, num_classes=3, normalize=None).to(torch.float32)

        torch.testing.assert_allclose(
            output,
            expected_output,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multiclass_confusion_matrix_with_probabilities(self) -> None:
        # should reduce to [2, 0, 1, 2, 1]
        input = torch.tensor(
            [
                [0.2948, 0.3343, 0.3709],
                [0.4988, 0.4836, 0.0176],
                [0.3727, 0.5145, 0.1128],
                [0.3759, 0.2115, 0.4126],
                [0.3076, 0.4226, 0.2698],
            ]
        )

        target = torch.tensor([0, 1, 2, 0, 0])

        expected_output = torch.tensor([[0, 1, 2], [1, 0, 0], [0, 1, 0]]).to(
            torch.float32
        )

        output = mccm(input, target, num_classes=3, normalize=None).to(torch.float32)

        torch.testing.assert_close(
            output,
            expected_output,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_sequential_prediction_model_with_ragged_inputs(torch_yoochoose_like, yoochoose_schema):
    input_module = tr.TabularSequenceFeatures.from_schema(
        yoochoose_schema,
        max_sequence_length=20,
        d_output=64,
        masking="causal",
    )
    prediction_task = tr.NextItemPredictionTask(weight_tying=True)
    transformer_config = tconf.XLNetConfig.build(
        d_model=64, n_head=8, n_layer=2, total_seq_length=20
    )
    model = transformer_config.to_torch_model(input_module, prediction_task)

    _ = model(torch_yoochoose_like)

    model.eval()

    inference_inputs = tr.data.tabular_sequence_testing_data.torch_synthetic_data(
        num_rows=10, min_session_length=1, max_session_length=4, ragged=True
    )
    inference_inputs_2 = tr.data.tabular_sequence_testing_data.torch_synthetic_data(
        num_rows=20, min_session_length=1, max_session_length=10, ragged=True
    )
    model_output = model(inference_inputs)

    # if the model is traced with ragged inputs it can only be called with ragged inputs
    # if the model is traced with padded inputs it can only be called with padded inputs
    traced_model = torch.jit.trace(model, inference_inputs, strict=False)
    traced_model_output = traced_model(inference_inputs)
    t.assert_close(model_output, traced_model_output)

    model_output = model(inference_inputs_2)
    traced_model_output = traced_model(inference_inputs_2)
    t.assert_close(model_output, traced_model_output)
//  [python]
def test_sequential_prediction_model_with_ragged_inputs(torch_yoochoose_like, yoochoose_schema):
    input_module = tr.TabularSequenceFeatures.from_schema(
        yoochoose_schema,
        max_sequence_length=20,
        d_output=64,
        masking="causal",
    )
    prediction_task = tr.NextItemPredictionTask(weight_tying=True)
    transformer_config = tconf.XLNetConfig.build(
        d_model=64, n_head=8, n_layer=2, total_seq_length=20
    )
    model = transformer_config.to_torch_model(input_module, prediction_task)

    _ = model(torch_yoochoose_like)

    model.eval()

    inference_inputs = tr.data.tabular_sequence_testing_data.torch_synthetic_data(
        num_rows=10, min_session_length=1, max_session_length=4, ragged=True
    )
    inference_inputs_2 = tr.data.tabular_sequence_testing_data.torch_synthetic_data(
        num_rows=20, min_session_length=1, max_session_length=10, ragged=True
    )
    model_output = model(inference_inputs)

    # if the model is traced with ragged inputs it can only be called with ragged inputs
    # if the model is traced with padded inputs it can only be called with padded inputs
    traced_model = torch.jit.trace(model, inference_inputs, strict=False)
    traced_model_output = traced_model(inference_inputs)
    torch.testing.assert_close(model_output, traced_model_output)

    model_output = model(inference_inputs_2)
    traced_model_output = traced_model(inference_inputs_2)
    torch.testing.assert_close(model_output, traced_model_output)
//  [python]
def assert_tensor_synced_across_pg(
    tensor: torch.Tensor,
    pg: dist.ProcessGroup,
    msg: Optional[Callable[[str], str]] = None,
    reference_rank: int = 0,
):
    """Assert that `tensor` is synced across `pg` with reference rank. Note that this always passes for reference rank"""
    if dist.get_rank(pg) == reference_rank:
        reference_tensor = tensor
    else:
        reference_tensor = torch.empty_like(tensor)
    dist.broadcast(
        reference_tensor,
        src=dist.get_global_rank(group=pg, group_rank=reference_rank),
        group=pg,
    )

    # TODO @nouamane: Getting Greatest absolute difference: 4.6e-10 at large scale when syncing tied weights
    torch.testing.assert_allclose(tensor, reference_tensor, msg=msg)
//  [python]
def assert_tensor_synced_across_pg(
    tensor: torch.Tensor,
    pg: dist.ProcessGroup,
    msg: Optional[Callable[[str], str]] = None,
    reference_rank: int = 0,
):
    """Assert that `tensor` is synced across `pg` with reference rank. Note that this always passes for reference rank"""
    if dist.get_rank(pg) == reference_rank:
        reference_tensor = tensor
    else:
        reference_tensor = torch.empty_like(tensor)
    dist.broadcast(
        reference_tensor,
        src=dist.get_global_rank(group=pg, group_rank=reference_rank),
        group=pg,
    )

    # TODO @nouamane: Getting Greatest absolute difference: 4.6e-10 at large scale when syncing tied weights
    torch.testing.assert_close(tensor, reference_tensor, msg=msg)
//  [python]
def _compare_models(
        self,
        m1: DistributedModelParallel,
        m2: DistributedModelParallel,
        is_deterministic: bool = True,
    ) -> None:
        sd1 = m1.state_dict()
        for key, value in m2.state_dict().items():
            v2 = sd1[key]
            if isinstance(value, ShardedTensor):
                assert isinstance(v2, ShardedTensor)
                self.assertEqual(len(value.local_shards()), len(v2.local_shards()))
                for dst, src in zip(value.local_shards(), v2.local_shards()):
                    if is_deterministic:
                        self.assertTrue(torch.equal(src.tensor, dst.tensor))
                    else:
                        rtol, atol = _get_default_rtol_and_atol(src.tensor, dst.tensor)
                        self.assertTrue(torch.allclose(src.tensor, dst.tensor, rtol=rtol, atol=atol))
            else:
                dst = value
                src = v2
                if is_deterministic:
                    self.assertTrue(torch.equal(src, dst))
                else:
                    rtol, atol = _get_default_rtol_and_atol(src, dst)
                    self.assertTrue(torch.allclose(src, dst, rtol=rtol, atol=atol))
//  [python]
def _compare_models(
        self,
        m1: DistributedModelParallel,
        m2: DistributedModelParallel,
        is_deterministic: bool = True,
    ) -> None:
        sd1 = m1.state_dict()
        for key, value in m2.state_dict().items():
            v2 = sd1[key]
            if isinstance(value, ShardedTensor):
                assert isinstance(v2, ShardedTensor)
                self.assertEqual(len(value.local_shards()), len(v2.local_shards()))
                for dst, src in zip(value.local_shards(), v2.local_shards()):
                    if is_deterministic:
                        self.assertTrue(torch.equal(src.tensor, dst.tensor))
                    else:
                        rtol, atol = _get_default_rtol_and_atol(src.tensor, dst.tensor)
                        torch.testing.assert_close(
                            src.tensor, dst.tensor, rtol=rtol, atol=atol
                        )
            else:
                dst = value
                src = v2
                if is_deterministic:
                    self.assertTrue(torch.equal(src, dst))
                else:
                    rtol, atol = _get_default_rtol_and_atol(src, dst)
                    torch.testing.assert_close(src, dst, rtol=rtol, atol=atol)
//  [python]
def before_tbi_sanity_checks(
    config: Config,
    parallel_context: ParallelContext,
    unwrapped_model: NanotronModel,
    grad_accumulator: GradientAccumulator,
) -> None:
    if not config.general.ignore_sanity_checks:
        # SANITY CHECK: Check that the model params are synchronized across dp
        for name, param in sorted(unwrapped_model.named_parameters(), key=lambda x: x[0]):
            assert_tensor_synced_across_pg(
                tensor=param,
                pg=parallel_context.dp_pg,
                msg=lambda err: f"{name} are not synchronized across DP {err}",
            )

        # SANITY CHECK: Tied weights are synchronized
        tied_params_list = sorted(
            get_tied_id_to_param(
                parameters=unwrapped_model.parameters(),
                root_module=unwrapped_model,
            ).items(),
            key=lambda x: x[0],
        )
        for (name, group_ranks), param in tied_params_list:
            group = parallel_context.world_ranks_to_pg[group_ranks]
            assert_tensor_synced_across_pg(
                tensor=param,
                pg=group,
                msg=lambda err: f"[Before train] Tied weights {name} are not synchronized. {err}",
            )

        # SANITY CHECK: Check that the grad accumulator buffers are ready for DDP
        if grad_accumulator is not None:
            for _, elt in grad_accumulator.fp32_grad_buffers.items():
                fp32_grad_buffer = elt["fp32_grad"]
                torch.allclose(
                    fp32_grad_buffer,
                    torch.zeros_like(fp32_grad_buffer),
                    atol=0,
                    rtol=0,
                    msg="Grad accumulator buffers must be zeroed in first accumulation step.",
                )

        # SANITY CHECK: run model specific sanity checks
        unwrapped_model.before_tbi_sanity_checks()
//  [python]
def before_tbi_sanity_checks(
    config: Config,
    parallel_context: ParallelContext,
    unwrapped_model: NanotronModel,
    grad_accumulator: GradientAccumulator,
) -> None:
    if not config.general.ignore_sanity_checks:
        # SANITY CHECK: Check that the model params are synchronized across dp
        for name, param in sorted(unwrapped_model.named_parameters(), key=lambda x: x[0]):
            assert_tensor_synced_across_pg(
                tensor=param,
                pg=parallel_context.dp_pg,
                msg=lambda err: f"{name} are not synchronized across DP {err}",
            )

        # SANITY CHECK: Tied weights are synchronized
        tied_params_list = sorted(
            get_tied_id_to_param(
                parameters=unwrapped_model.parameters(),
                root_module=unwrapped_model,
            ).items(),
            key=lambda x: x[0],
        )
        for (name, group_ranks), param in tied_params_list:
            group = parallel_context.world_ranks_to_pg[group_ranks]
            assert_tensor_synced_across_pg(
                tensor=param,
                pg=group,
                msg=lambda err: f"[Before train] Tied weights {name} are not synchronized. {err}",
            )

        # SANITY CHECK: Check that the grad accumulator buffers are ready for DDP
        if grad_accumulator is not None:
            for _, elt in grad_accumulator.fp32_grad_buffers.items():
                fp32_grad_buffer = elt["fp32_grad"]
                torch.testing.assert_close(
                    fp32_grad_buffer,
                    torch.zeros_like(fp32_grad_buffer),
                    atol=0,
                    rtol=0,
                    msg="Grad accumulator buffers must be zeroed in first accumulation step.",
                )

        # SANITY CHECK: run model specific sanity checks
        unwrapped_model.before_tbi_sanity_checks()
//  [python]
def testLargeTensor(self):
        t = torch.zeros(2359332864, dtype=torch.half, device='cuda')
        t2 = torch.zeros(2359332864, dtype=torch.half, device='cuda')
        grad = torch.randn_like(t)
        t.grad = grad
        t2.grad = grad
        params = [t]
        params2 = [t2]
        optimizer = apex.optimizers.FusedAdam(params, lr=self.lr)
        optimizer.step()
        optimizer2 = torch.optim.Adam(params2, lr=self.lr)
        torch.testing.assert_allclose(t, t2)
        torch.cuda.synchronize()
//  [python]
def testLargeTensor(self):
        t = torch.zeros(2359332864, dtype=torch.half, device='cuda')
        t2 = torch.zeros(2359332864, dtype=torch.half, device='cuda')
        grad = torch.randn_like(t)
        t.grad = grad
        t2.grad = grad
        params = [t]
        params2 = [t2]
        optimizer = apex.optimizers.FusedAdam(params, lr=self.lr)
        optimizer.step()
        optimizer2 = torch.optim.Adam(params2, lr=self.lr)
        torch.testing.assert_close(t, t2)
        torch.cuda.synchronize()
//  [python]
print("\nParameters of each stage:")
for stage_idx in range(pipe.num_stages):
    print(f"\nStage {stage_idx}:")
    stage_mod = pipe.get_stage_module(stage_idx)
    for param_name, param in stage_mod.named_parameters():
        assert (
            param_name in orig_state_dict
        ), f"{param_name} not in original state dict"
        print(f"{param_name}: {param.size()}")

# Check equivalence
ref = mod(x, constant)
out = pipe(x, constant)[0]
t.assert_close(out, ref)
//  [python]
print("\nParameters of each stage:")
for stage_idx in range(pipe.num_stages):
    print(f"\nStage {stage_idx}:")
    stage_mod = pipe.get_stage_module(stage_idx)
    for param_name, param in stage_mod.named_parameters():
        assert (
            param_name in orig_state_dict
        ), f"{param_name} not in original state dict"
        print(f"{param_name}: {param.size()}")

# Check equivalence
ref = mod(x, constant)
out = pipe(x, constant)[0]
torch.testing.assert_close(out, ref)
//  [python]
def testGradScaler(self):
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=False)
        scaler = torch.cuda.amp.GradScaler(enabled=True)
        scaler_ = torch.cuda.amp.GradScaler(enabled=True)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model(x)
                loss = ((gt - y) ** 2).mean()

            scaler.scale(loss).backward()
            scaler.step(self.optimizer)
            scaler.update()
            
            # DUT
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model_(x)
                loss_ = ((gt_ - y) ** 2).mean()

            scaler_.scale(loss_).backward()
            scaler_.step(optimizer_)
            scaler_.update()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_allclose(m.weight, m_.weight, atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_allclose(m.weight.grad, m_.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()

            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def testGradScaler(self):
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=False)
        scaler = torch.cuda.amp.GradScaler(enabled=True)
        scaler_ = torch.cuda.amp.GradScaler(enabled=True)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model(x)
                loss = ((gt - y) ** 2).mean()

            scaler.scale(loss).backward()
            scaler.step(self.optimizer)
            scaler.update()
            
            # DUT
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model_(x)
                loss_ = ((gt_ - y) ** 2).mean()

            scaler_.scale(loss_).backward()
            scaler_.step(optimizer_)
            scaler_.update()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_close(m.weight, m_.weight, atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_close(m.weight.grad, m_.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()

            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def testGradScalerCapturable(self):
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=True)
        scaler = torch.cuda.amp.GradScaler(enabled=True)
        scaler_ = torch.cuda.amp.GradScaler(enabled=True)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model(x)
                loss = ((gt - y) ** 2).mean()

            scaler.scale(loss).backward()
            scaler.step(self.optimizer)
            scaler.update()
            
            # DUT
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model_(x)
                loss_ = ((gt_ - y) ** 2).mean()

            scaler_.scale(loss_).backward()
            scaler_.step(optimizer_)
            scaler_.update()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_allclose(m.weight, m_.weight, atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_allclose(m.weight.grad, m_.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()

            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def testGradScalerCapturable(self):
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=True)
        scaler = torch.cuda.amp.GradScaler(enabled=True)
        scaler_ = torch.cuda.amp.GradScaler(enabled=True)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model(x)
                loss = ((gt - y) ** 2).mean()

            scaler.scale(loss).backward()
            scaler.step(self.optimizer)
            scaler.update()
            
            # DUT
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model_(x)
                loss_ = ((gt_ - y) ** 2).mean()

            scaler_.scale(loss_).backward()
            scaler_.step(optimizer_)
            scaler_.update()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_close(m.weight, m_.weight, atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_close(m.weight.grad, m_.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()

            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def testNative(self):
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=False)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            y = self.model(x)
            loss = ((gt - y) ** 2).mean()

            loss.backward()
            self.optimizer.step()
            
            # DUT
            y = self.model_(x)
            loss_ = ((gt_ - y) ** 2).mean()

            loss_.backward()
            optimizer_.step()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_allclose(m.weight, m_.weight, atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_allclose(m.weight.grad, m_.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()
            
            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def testNative(self):
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=False)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            y = self.model(x)
            loss = ((gt - y) ** 2).mean()

            loss.backward()
            self.optimizer.step()
            
            # DUT
            y = self.model_(x)
            loss_ = ((gt_ - y) ** 2).mean()

            loss_.backward()
            optimizer_.step()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_close(m.weight, m_.weight, atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_close(m.weight.grad, m_.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()
            
            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def test_optimize_runtime_with_given_memory(max_memory, optimal_soln):
    data = [
        ("aten.copy_", 5, 0),
        ("aten.add", 5, 100),
        ("aten.div", 8, 100),
        ("aten.mm", 15, 120),
        ("aten.native_dropout", 15, 0),
        ("aten.linear", 9, 100),
        ("aten.t", 1, 0),
        ("aten.relu_", 5, 0),
    ]

    inplace_ops = [(0, 0), (7, 5)]
    view_like_ops = [6]
    rand_ops = [4]

    runtimes = torch.tensor([x[1] for x in data], dtype=torch.float64)
    memory = torch.tensor([x[2] for x in data], dtype=torch.float64)

    out = _optimize_runtime_with_given_memory(
        memory,
        runtimes,
        max_memory,
        view_like_ops,
        inplace_ops,
        rand_ops,
        force_store_random=False,
    )
    t.assert_close(optimal_soln, out)
//  [python]
def test_optimize_runtime_with_given_memory(max_memory, optimal_soln):
    data = [
        ("aten.copy_", 5, 0),
        ("aten.add", 5, 100),
        ("aten.div", 8, 100),
        ("aten.mm", 15, 120),
        ("aten.native_dropout", 15, 0),
        ("aten.linear", 9, 100),
        ("aten.t", 1, 0),
        ("aten.relu_", 5, 0),
    ]

    inplace_ops = [(0, 0), (7, 5)]
    view_like_ops = [6]
    rand_ops = [4]

    runtimes = torch.tensor([x[1] for x in data], dtype=torch.float64)
    memory = torch.tensor([x[2] for x in data], dtype=torch.float64)

    out = _optimize_runtime_with_given_memory(
        memory,
        runtimes,
        max_memory,
        view_like_ops,
        inplace_ops,
        rand_ops,
        force_store_random=False,
    )
    torch.testing.assert_close(optimal_soln, out)
//  [python]
def testGradScalerCapturableMaster(self):
        # Cast conv layers to FP16
        for m in self.model_.modules():
            if m.__class__ in [torch.nn.Conv2d]:
                m.half()
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=True, master_weights=True)
        scaler = torch.cuda.amp.GradScaler(enabled=True)
        scaler_ = torch.cuda.amp.GradScaler(enabled=True)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model(x)
                loss = ((gt - y) ** 2).mean()

            scaler.scale(loss).backward()
            scaler.step(self.optimizer)
            scaler.update()

            # DUT
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model_(x)
                loss_ = ((gt_ - y) ** 2).mean()

            scaler_.scale(loss_).backward()
            scaler_.step(optimizer_)
            scaler_.update()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_allclose(m.weight, m_.weight.float(), atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_allclose(m.weight.grad, m_.weight.grad.float(), atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()

            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def testGradScalerCapturableMaster(self):
        # Cast conv layers to FP16
        for m in self.model_.modules():
            if m.__class__ in [torch.nn.Conv2d]:
                m.half()
        params_ = [p for p in self.model_.parameters() if p.requires_grad]
        optimizer_ = apex.optimizers.FusedAdam(params_, lr=self.lr, capturable=True, master_weights=True)
        scaler = torch.cuda.amp.GradScaler(enabled=True)
        scaler_ = torch.cuda.amp.GradScaler(enabled=True)

        for i in range(100):
            x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last)
            x_ = x.clone()
            gt = torch.rand([32, 10]).cuda()
            gt_ = gt.clone()

            # Reference
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model(x)
                loss = ((gt - y) ** 2).mean()

            scaler.scale(loss).backward()
            scaler.step(self.optimizer)
            scaler.update()

            # DUT
            with torch.cuda.amp.autocast(enabled=True):
                y = self.model_(x)
                loss_ = ((gt_ - y) ** 2).mean()

            scaler_.scale(loss_).backward()
            scaler_.step(optimizer_)
            scaler_.update()

            for module in zip(self.model.modules(), self.model_.modules()):
                m = module[0]
                m_ = module[1]
                if isinstance(m, nn.Conv2d) or isinstance(m_, nn.Linear):
                    torch.testing.assert_close(m.weight, m_.weight.float(), atol=1e-3, rtol=1e-3, equal_nan=True)
                    torch.testing.assert_close(m.weight.grad, m_.weight.grad.float(), atol=1e-3, rtol=1e-3, equal_nan=True)

            # Init for next iteration
            self.optimizer.zero_grad()
            optimizer_.zero_grad()

            self.model_.load_state_dict(copy.deepcopy(self.model.state_dict()))
//  [python]
def test_optimal_checkpoint_policy(
    device, memory_budget, inplace, random, first_inplace
):
    if first_inplace and inplace:
        pytest.skip("This case is degenerate and doesn't work with vanilla PyTorch")
    torch.manual_seed(42)
    dtype = torch.float16
    modules = _get_model_blocks(
        3, dtype, device, inplace=inplace, random=random, first_inplace=first_inplace
    )
    inputs = torch.rand(32, 128, 10, dtype=dtype, device=device)

    policy_fn = get_optimal_checkpoint_policy(
        modules[0], inputs, memory_budget=memory_budget
    )
    model = _Model(modules, policy_fn)
    model_ref = torch.nn.Sequential(*deepcopy(modules))

    grad = torch.rand_like(inputs)

    torch.manual_seed(42)
    out = model(inputs.clone())
    out.backward(grad)

    torch.manual_seed(42)
    out_ref = model_ref(inputs.clone())
    out_ref.backward(grad)

    t.assert_close(out, out_ref)

    for p, p_ref in zip(model.parameters(), model_ref.parameters()):
        t.assert_close(p.grad, p_ref.grad)
//  [python]
def test_optimal_checkpoint_policy(
    device, memory_budget, inplace, random, first_inplace
):
    if first_inplace and inplace:
        pytest.skip("This case is degenerate and doesn't work with vanilla PyTorch")
    torch.manual_seed(42)
    dtype = torch.float16
    modules = _get_model_blocks(
        3, dtype, device, inplace=inplace, random=random, first_inplace=first_inplace
    )
    inputs = torch.rand(32, 128, 10, dtype=dtype, device=device)

    policy_fn = get_optimal_checkpoint_policy(
        modules[0], inputs, memory_budget=memory_budget
    )
    model = _Model(modules, policy_fn)
    model_ref = torch.nn.Sequential(*deepcopy(modules))

    grad = torch.rand_like(inputs)

    torch.manual_seed(42)
    out = model(inputs.clone())
    out.backward(grad)

    torch.manual_seed(42)
    out_ref = model_ref(inputs.clone())
    out_ref.backward(grad)

    torch.testing.assert_close(out, out_ref)

    for p, p_ref in zip(model.parameters(), model_ref.parameters()):
        torch.testing.assert_close(p.grad, p_ref.grad)
//  [python]
def test_split_module_dead_code(self):
        class ModWithDeadCode(torch.nn.Module):
            def forward(self, x):
                output = x * 2  # we want this
                dead_line = x + 2  # this is dead
                return output

        mod = ModWithDeadCode()
        traced = torch.fx.symbolic_trace(mod)

        # split into before (0), target (1), and after(2)
        saw_mul = False

        def split_callback(n):
            nonlocal saw_mul
            if n.target == operator.mul:
                saw_mul = True
                return 1

            if not saw_mul:
                return 0
            if saw_mul:
                return 2

        split = split_module(traced, mod, split_callback)

        x = torch.randn((5,))
        t.assert_close(
            split(x), traced(x)
        )
//  [python]
def test_split_module_dead_code(self):
        class ModWithDeadCode(torch.nn.Module):
            def forward(self, x):
                output = x * 2  # we want this
                dead_line = x + 2  # this is dead
                return output

        mod = ModWithDeadCode()
        traced = torch.fx.symbolic_trace(mod)

        # split into before (0), target (1), and after(2)
        saw_mul = False

        def split_callback(n):
            nonlocal saw_mul
            if n.target == operator.mul:
                saw_mul = True
                return 1

            if not saw_mul:
                return 0
            if saw_mul:
                return 2

        split = split_module(traced, mod, split_callback)

        x = torch.randn((5,))
        torch.testing.assert_close(
            split(x), traced(x)
        )
//  [python]
def test_split_module_kwargs_expansion(self):
        class ModuleWithKwargsExpansion(torch.nn.Module):
            def forward(self, x, **kwargs):
                return x + kwargs['foo']

        mod = ModuleWithKwargsExpansion()
        traced = torch.fx.symbolic_trace(mod)

        seen_getitem = False

        def split_callback(n):
            nonlocal seen_getitem
            split_idx = int(seen_getitem)
            if n.target == operator.getitem:
                seen_getitem = True
            return split_idx

        split = split_module(traced, mod, split_callback)

        x = torch.randn(5, 3)
        foo = torch.randn(5, 3)
        t.assert_close(split(x, foo=foo), traced(x, foo=foo))
//  [python]
def test_split_module_kwargs_expansion(self):
        class ModuleWithKwargsExpansion(torch.nn.Module):
            def forward(self, x, **kwargs):
                return x + kwargs['foo']

        mod = ModuleWithKwargsExpansion()
        traced = torch.fx.symbolic_trace(mod)

        seen_getitem = False

        def split_callback(n):
            nonlocal seen_getitem
            split_idx = int(seen_getitem)
            if n.target == operator.getitem:
                seen_getitem = True
            return split_idx

        split = split_module(traced, mod, split_callback)

        x = torch.randn(5, 3)
        foo = torch.randn(5, 3)
        torch.testing.assert_close(split(x, foo=foo), traced(x, foo=foo))
//  [python]
def test_split_module_default_arg(self):
        class ModelToTrace(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.lin = torch.nn.Linear(512, 512)

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

                if targets is not None:
                    x = x + targets

                return x

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

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

        x = torch.randn(50, 512)
        t.assert_close(split(x), traced(x))