Diff Coverage

    is_compiling = getattr(plat, "is_compiling", None)
    checkpoint_compile = callable(is_compiling) and is_compiling() is True  # pylint: disable=not-callable
    if checkpoint_compile:
        unsupported = []
        if swap_inputs:
            unsupported.append("swap_inputs")
        if group_swap:
            unsupported.append("group_swap")
        if context_fn is not None:
            unsupported.append("custom context_fn")
        if kwargs.get("use_reentrant", False):
            unsupported.append("use_reentrant=True")
        if unsupported:
            raise ValueError(
                "HyperParallel checkpoint compile mode does not support: "
                + ", ".join(unsupported)
                + ". Use Torch-native non-reentrant checkpointing with optional "
                "SAVE/RECOMPUTE selective policies."

                + ", ".join(unsupported)
                + ". Use Torch-native non-reentrant checkpointing with optional "
                "SAVE/RECOMPUTE selective policies."
            )
        composed_context_fn = (
            partial(plat.create_native_selective_checkpoint_contexts, policy_fn)
            if policy_fn is not None
            else None
        )

        >>> output = swap(layer, x, policy_fn=lambda t: CheckpointPolicy.MUST_SAVE)
    """
    is_compiling = getattr(plat, "is_compiling", None)
    if callable(is_compiling) and is_compiling() is True:  # pylint: disable=not-callable
        raise ValueError(
            "HyperParallel activation swap is not supported in compile mode. "
            "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
        )
    with plat.async_save_on_cpu(policy_fn=policy_fn, group_swap=group_swap):

    @staticmethod
    def create_native_selective_checkpoint_contexts(policy_fn: Callable) -> Any:
        """Create framework-native selective checkpoint contexts for compile."""
        raise NotImplementedError(
            "Native selective checkpoint compile is not supported by this platform"
        )

    @staticmethod

def _raise_if_compiling(feature_name: str) -> None:
    """Reject HyperParallel activation-swap extensions during Torch capture."""
    if torch.compiler.is_compiling():
        raise ValueError(
            f"HyperParallel {feature_name} is not supported in compile mode. "
            "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
        )

# limitations under the License.
# ============================================================================
"""Adapters for using PyTorch-native selective checkpointing under compile."""

from functools import partial
from typing import Any, Callable, Tuple

from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy


_SUPPORTED_POLICY_NAMES = (
    "MUST_SAVE",
    "PREFER_SAVE",
    "MUST_RECOMPUTE",
    "PREFER_RECOMPUTE",

    "PREFER_RECOMPUTE",
)


def _to_torch_checkpoint_policy(policy: Any) -> Any:
    """Convert supported HyperParallel policies to the native Torch enum."""
    from torch.utils import checkpoint as torch_checkpoint  # pylint: disable=C0415

    torch_policy_cls = torch_checkpoint.CheckpointPolicy
    supported_native = {
        getattr(torch_policy_cls, name) for name in _SUPPORTED_POLICY_NAMES
    }
    if isinstance(policy, torch_policy_cls):
        if policy in supported_native:
            return policy
        raise ValueError(
            f"Torch checkpoint policy {policy.name} is not supported by "
            "HyperParallel compile mode. Only SAVE and RECOMPUTE policies are supported."
        )
    if isinstance(policy, CheckpointPolicy):
        if policy.name in _SUPPORTED_POLICY_NAMES:
            return getattr(torch_policy_cls, policy.name)
        raise ValueError(
            f"HyperParallel checkpoint policy {policy.name} is not supported in compile mode. "
            "Only SAVE and RECOMPUTE policies are supported."
        )
    raise TypeError(
        "Selective checkpoint policy_fn must return a HyperParallel or Torch "
        f"CheckpointPolicy, but got {type(policy).__name__}."
    )

        f"CheckpointPolicy, but got {type(policy).__name__}."
    )


def _torch_policy_adapter(
    policy_fn: Callable, torch_context: Any, op: Any, *args: Any, **kwargs: Any
) -> Any:
    """Pass native Torch inputs through and adapt only the policy result."""
    return _to_torch_checkpoint_policy(policy_fn(torch_context, op, *args, **kwargs))


def create_native_selective_checkpoint_contexts(policy_fn: Callable) -> Tuple[Any, Any]:
    """Create Torch-native selective-checkpoint contexts for compile capture."""
    if not callable(policy_fn):
        raise TypeError("policy_fn must be callable in HyperParallel compile mode.")
    from torch.utils import checkpoint as torch_checkpoint  # pylint: disable=C0415

    return torch_checkpoint.create_selective_checkpoint_contexts(
        partial(_torch_policy_adapter, policy_fn)
    )


__all__ = ["create_native_selective_checkpoint_contexts"]

    @_dynamo_disable
    def _grouped_forward_pre_hook(self, cell, args, kwargs):
        """Run the grouped FSDP pre-forward hook outside Dynamo tracing."""
        return super()._grouped_forward_pre_hook(cell, args, kwargs)

    def _register_forward_module_hook(self, mod, hook) -> None:
        """Register forward hook; use ``always_call=True`` when supported (matches PyTorch FSDP)."""
        sig = inspect.signature(mod.register_forward_hook)

def _ensure_contiguous(x):
    """Return a contiguous copy of *x* if not already contiguous."""
    if torch.compiler.is_compiling():
        return x.contiguous()
    if not x.is_contiguous() or x.storage_offset() != 0:
        return x.contiguous()
    return x

    @staticmethod
    def create_native_selective_checkpoint_contexts(policy_fn: Callable) -> Any:
        """Create Torch-native selective-checkpoint contexts for compile."""
        # pylint: disable=C0415
        from hyper_parallel.platform.torch.activation_checkpoint.native_compile import (
            create_native_selective_checkpoint_contexts,
        )
        return create_native_selective_checkpoint_contexts(policy_fn)

    @staticmethod
    def async_save_on_cpu(policy_fn=None, group_swap: bool = False):
        # pylint: disable=C0415