class ShardMoESparseExpertsParallel(ParallelStyle):
    """
    Parallel style that shards MoE experts across a specific mesh dimension.

    This style is designed for ``MoELayer`` instances using ``GroupedLinear`` for experts.
    It splits the experts across the specified
    dimension of the device mesh (Expert Parallelism). Other dimensions in the
    mesh treat the parameters as Replicated.

    It also initializes the necessary distributed communication groups within the
    MoE layer to handle token dispatching.
    """

    def __init__(self, shard_dim_name: str):
        self._shard_dim_name = shard_dim_name

    def _partition_experts(self, module_name: str, mod: nn.Module, device_mesh: DeviceMesh):
        if not isinstance(mod, GroupedLinear):
            raise TypeError("This plan should be applied only on GroupedLinear")

        mesh_dim_names = device_mesh.mesh_dim_names

        if mesh_dim_names is None:
            raise ValueError("This plan should be applied only on named DeviceMeshes")

        placements = [
            Shard(0) if dim_name == self._shard_dim_name else Replicate()
            for dim_name
            in mesh_dim_names
        ]
        weight = nn.Parameter(
            distribute_tensor(mod.weight, device_mesh, placements),
            requires_grad=mod.weight.requires_grad
        )
        mod.weight = weight

    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
        if not isinstance(module, MoELayer):
            raise TypeError("This plan should be applied only on MoELayer")

        module.enable_distributed_communicator(device_mesh.get_group(self._shard_dim_name))

        for submod in module.modules():
            if isinstance(submod, GroupedLinear):
                distribute_module(submod, device_mesh, self._partition_experts)

        return module

class ToLocalParallel(ParallelStyle):
    """
    Parallel style that distributes parameters and gradients but executes with local tensors.

    This style wraps standard tensor distribution (via ``DTensor``) but injects
    runtime hooks to temporarily unwrap ``DTensor`` parameters into local ``torch.Tensor``
    during the forward pass.

    This is useful for parallel strategies (like Replicate)
    where the underlying calculation logic is not DTensor-aware, but the parameters must remain
    distributed for gradient synchronization and for distributed checkpointing.
    """

    def __init__(self, param_placement: tuple[Placement, ...], grad_placement: tuple[Placement, ...]):
        """
        Constructs ToLocalParallel object.

        Args:
            param_placement: Tuple of placements defining how parameters are distributed.
            grad_placement: Tuple of placements defining how gradients are synchronized.
        """

        self._grad_placement = grad_placement
        self._param_placement = param_placement

    def _distribute_params(self, name: str, module: nn.Module, device_mesh: DeviceMesh):
        for param_name, param in module.named_parameters(recurse=False):
            new_param = nn.Parameter(
                distribute_tensor(param.data, device_mesh, self._param_placement),
                requires_grad=param.requires_grad
            )

            module.register_parameter(param_name, new_param)

    def _apply(self, master_module: nn.Module, device_mesh: DeviceMesh):
        patched_classes = {}
        original_classes = {}

        for submod_name, submod in master_module.named_modules():
            param_names = [name for name, p in submod.named_parameters(recurse=False)]
            patched_classes[submod_name] = _build_to_local_patched_class(submod, self._grad_placement, param_names)
            original_classes[submod_name] = submod.__class__

            distribute_module(
                submod,
                device_mesh,
                self._distribute_params
            )

        master_module.register_forward_pre_hook(_ModulePatch(patched_classes))
        master_module.register_forward_hook(_ModulePatch(original_classes))