Diff Coverage

    @staticmethod
    def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
        """Override in subclass to define the forward computation on local tensors."""
        raise NotImplementedError("Subclasses must implement forward()")

    @staticmethod
    def backward(ctx: Any, *grad_outputs: Any) -> Any:
        """Override in subclass to define the backward computation on local tensors."""
        raise NotImplementedError("Subclasses must implement backward()")

    @classmethod
    def apply(cls, *args: Any, **kwargs: Any) -> Any:
        """Execute the function, routing to distributed dispatch when DTensor inputs are present.

                raise ValueError(
                    f"{cls.__name__} received DTensor inputs but '_op_name' is not set. "
                    "Set '_op_name' on the subclass and register a matching DistributedOp."
                )
            return _OP_DISPATCHER.dispatch(cls._get_local_callable(), args, kwargs)
        return super().apply(*args, **kwargs)

    @classmethod
    def _get_local_callable(cls) -> _LocalCallable:

        if '_local_callable' not in cls.__dict__:
            def _local_fn(*a, **kw):
                # Called by _OP_DISPATCHER with extracted local (non-DTensor) tensors.
                # has_dtensor will be False here, so super().apply() is taken directly.
                return cls.apply(*a, **kw)

            cls._local_callable = _LocalCallable(_local_fn, cls._op_name)
        return cls._local_callable