Diff Coverage

    def hooked_leaf_step(optimizer: Any, *args: Any, **kwargs: Any) -> Any:
        """Dispatch one raw optimizer leaf through its swap wrapper."""
        current = wrapper_ref()
        if current is None:
            raise RuntimeError("Swap optimizer wrapper has been released.")
        if optimizer is not current.optimizer:
            raise ValueError("Optimizer step hooks cannot replace optimizer identity.")
        return current._step_impl(*args, **kwargs)

    return torch.optim.Optimizer.profile_hook_step(hooked_leaf_step)

        """Execute the swap pipeline between PyTorch pre- and post-hooks."""
        # ``step()`` already rejected the caller's closure; whatever arrives here
        # is either the forwarded ``closure=None`` or a pre-hook rewrite.
        if len(args) > 1 or set(kwargs) - {"closure"}:
            raise ValueError("Swap optimizer does not support extra step arguments.")
        if args and "closure" in kwargs:
            raise TypeError("closure was passed both positionally and by keyword.")
        closure = args[0] if args else kwargs.get("closure")
        if closure is not None:
            raise ValueError("Swap optimizer does not support closure.")

        with self._no_grad_context():
            step_context = self.adapter.prepare_step()
            units = self.adapter.iter_update_units(step_context)