Diff Coverage

        if swap_inputs:
            unsupported.append("swap_inputs")
        if group_swap:
            unsupported.append("group_swap")
        if cpu_pool is not None:
            unsupported.append("cpu_pool")
        if context_fn is not None:
            unsupported.append("custom context_fn")
        if kwargs.get("use_reentrant", False):
            unsupported.append("use_reentrant=True")

        factories: list = [create_recompute_contexts]
        if policy_fn is not None:
            selective_kwargs = {"group_swap": group_swap}
            if cpu_pool is not None:
                selective_kwargs["cpu_pool"] = cpu_pool
            factories.append(partial(plat.create_selective_checkpoint_contexts, policy_fn, **selective_kwargs))
        if context_fn is not None:
            factories.append(context_fn)

    if swap_inputs:
        async_kwargs = {"group_swap": group_swap}
        if cpu_pool is not None:
            async_kwargs["cpu_pool"] = cpu_pool
        context = partial(plat.async_save_on_cpu, **async_kwargs)
    else:
        context = contextlib.nullcontext
    with context():

            "Use Torch-native non-reentrant checkpointing with SAVE/RECOMPUTE policies."
        )
    async_kwargs = {"policy_fn": policy_fn, "group_swap": group_swap}
    if cpu_pool is not None:
        async_kwargs["cpu_pool"] = cpu_pool
    with plat.async_save_on_cpu(**async_kwargs):
        return function(*args, **kwargs)

        # event objects so each runtime event is queried only once per scan.
        event_results: Dict[int, Tuple[Any, bool]] = {}
        for capacity in list(self._pending):
            if capacity < minimum_capacity:
                continue
            still_pending = []
            for block in self._pending[capacity]:
                event = block.event
                event_key = id(event)

        """Associate a checked-out view with its owning allocation."""
        key = _storage_key(view)
        owner = self._blocks.get(key)
        if owner is not None and owner is not block:
            raise RuntimeError("Two host buffers exposed the same storage identity.")
        old_key = block.storage_key
        if old_key is not None and old_key != key and self._blocks.get(old_key) is block:
            del self._blocks[old_key]
        self._blocks[key] = block

            block = self._find_available_locked(aligned_size)
            if block is not None:
                try:
                    return self._checkout_locked(block, size)
                except Exception:
                    block.state = _AVAILABLE
                    self._available[block.capacity].append(block)
                    raise

            if self._total_allocated + aligned_size <= self._max_host_bytes:
                self._total_allocated += aligned_size
                reserved = True

            block = _Block(buffer, aligned_size, _IN_USE)
            try:
                with self._lock:
                    view = self._checkout_locked(block, size)
            except Exception:
                with self._lock:
                    self._total_allocated -= aligned_size
                raise
            return view

        event = block.event
        try:

        event = block.event
        try:
            event.synchronize()
        except Exception:
            with self._lock:
                block.state = _PENDING
                block.event = event
                self._pending[block.capacity].insert(0, block)
            raise
        block.event = None
        try:
            with self._lock:
                return self._checkout_locked(block, size)
        except Exception:
            with self._lock:
                block.state = _AVAILABLE
                self._available[block.capacity].append(block)
            raise

    def release(self, tensor: Any, event: Optional[Any] = None) -> None:
        """Return an acquired view to the pool, optionally after an async event."""
        key = _storage_key(tensor)

    def release_cpu_buffer(self, event=None):
        """Release an explicitly pooled host buffer exactly once."""
        if self.cpu_pool is None or self._cpu_pool_buffer is None:
            return
        release_tensor = self.val_cpu if self.val_cpu is not None else self._cpu_pool_buffer
        self.cpu_pool.release(release_tensor, event=event)
        self._cpu_pool_buffer = None
        self.val_cpu = None

    def wait_load(self, release_event=None):
        """change state to device after async load is done"""
        if self._state == self.STATE_NON_TENSOR or self._keep_on_device or self._duplicate_swap:

                self.val_cpu = platform.empty_like(
                    self.val, device="cpu", pin_memory=True
                )
            else:
                logical_bytes = self.val.numel() * platform.get_element_size(self.val)
                self._cpu_pool_buffer = self.cpu_pool.acquire(logical_bytes)
                try:
                    self.val_cpu = self._cpu_pool_buffer.view(self.val.dtype).reshape(self.val.shape)
                except Exception:
                    self.release_cpu_buffer()
                    raise
        try:
            if self.cpu_pool is not None or self.is_slice_tensor:
                self.val_cpu.copy_(self.val, non_blocking=True)
            else:
                self.val_cpu.untyped_storage().copy_(self.val.untyped_storage(), non_blocking=True)
        except Exception:
            if self.cpu_pool is not None and self._cpu_pool_buffer is not None:
                release_event = platform.new_event()
                release_event.record(platform.get_current_stream())
                self.release_cpu_buffer(release_event)
            self.val_cpu = None
            raise
        self._state = self.STATE_D2H

    def wait_offload(self):
        """wait offload to host and free device memory"""

        self.clear()

    def release_cpu_buffers(self, event=None):
        """Release all explicitly pooled host buffers held by this storage."""
        def _release(x):
            if isinstance(x, SwapTensor):
                x.release_cpu_buffer(event=event)
            return x

        for storage_list in self.values():
            for item in storage_list:
                platform.tree_map(_release, item)

    def wait_offload(self):
        """wait offload for all tensors in swap storage"""
        def _wait_offload(x):

                        cpu_pool = bucket["cpu_pool"]
                        if cpu_pool is None:
                            cpu_buf = _get_cpu_pinned_buf(dtype_key, numel, bucket["dtype"])
                        else:
                            raw_buf = cpu_pool.acquire(bucket["total_bytes"])
                            try:
                                cpu_buf = raw_buf.view(bucket["dtype"])
                            except Exception:
                                cpu_pool.release(raw_buf)
                                raise
                        group_cpu_bufs[bucket_key] = cpu_buf
                        cpu_buf[:numel].copy_(group_device_bufs[bucket_key], non_blocking=True)
                except Exception:
                    release_event = platform.new_event()
                    release_event.record(copy_stream)
                    for bucket_key, cpu_buf in group_cpu_bufs.items():
                        bucket = self._packed_buckets[bucket_key]
                        if bucket["cpu_pool"] is not None:
                            bucket["cpu_pool"].release(cpu_buf, event=release_event)
                        else:
                            _return_cpu_pinned_buf(cpu_buf)
                    raise
                self._group_device_buf = group_device_bufs
                self._group_cpu_buf = group_cpu_bufs

            # Slice tensors use the existing per-tensor path.

    def release_cpu_buffers(self, event=None):
        """Release staging buffers immediately or defer until ``event`` completes."""
        if self._group_cpu_buf is None:
            return
        for bucket_key, buf in self._group_cpu_buf.items():
            bucket = self._packed_buckets.get(bucket_key)
            if bucket is not None and bucket["cpu_pool"] is not None:
                bucket["cpu_pool"].release(buf, event=event)
            else:
                _return_cpu_pinned_buf(buf)
        self._group_cpu_buf = None

    def wait_load(self):
        """Wait for load to complete for all storages in the group.

        for event in (group._offload_event, group._load_event):
            if event is not None:
                event.synchronize()
        for storage in group._storages:
            storage.release_cpu_buffers()
        group.release_cpu_buffers()
        group._storages.clear()

    def get_current_group_name(self) -> str:

    @staticmethod
    def swap_wrapper(module, policy_fn=None, group_swap=False, cpu_pool=None):
        # pylint: disable=C0415
        from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_wrapper
        return swap_wrapper(module, policy_fn=policy_fn, group_swap=group_swap, cpu_pool=cpu_pool)

    @staticmethod
    def swap_tensor_wrapper(target, tag=None, group_swap=False, cpu_pool=None):
        # pylint: disable=C0415

    @staticmethod
    def swap_tensor_wrapper(target, tag=None, group_swap=False, cpu_pool=None):
        # pylint: disable=C0415
        from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import swap_tensor_wrapper
        return swap_tensor_wrapper(target, tag=tag, group_swap=group_swap, cpu_pool=cpu_pool)

    @staticmethod
    def get_class_activation_wrapper():
        # pylint: disable=C0415

        policy_fn_or_list, allow_cache_entry_mutation=False, group_swap=False, cpu_pool=None
    ):
        # pylint: disable=C0415
        from hyper_parallel.platform.torch.activation_checkpoint.sac import create_selective_checkpoint_contexts
        return create_selective_checkpoint_contexts(
            policy_fn_or_list, allow_cache_entry_mutation, group_swap, cpu_pool
        )

    @staticmethod

    @staticmethod
    def async_save_on_cpu(policy_fn=None, group_swap: bool = False, cpu_pool=None):
        # pylint: disable=C0415
        from hyper_parallel.platform.torch.activation_checkpoint.activation_swap import AsyncSaveOnCpu
        return AsyncSaveOnCpu(policy_fn, group_swap=group_swap, cpu_pool=cpu_pool)

    @staticmethod
    def get_element_size(tensor):
        """Get Tensor Element Size"""