Diff Coverage

            def _convert(val):
                if hasattr(val, "to_dict"):
                    return val.to_dict()
                if dataclasses.is_dataclass(val):
                    return dataclasses.asdict(val)
                if isinstance(val, (list, tuple)):
                    return type(val)(_convert(v) for v in val)
                if isinstance(val, dict):
                    return {k: _convert(v) for k, v in val.items()}
                if isinstance(val, (str, int, float, bool, type(None))):
                    return val
                if callable(val):
                    return repr(val)
                logger.warning(
                    "Config field value of type %s may not be JSON serializable",
                    type(val).__name__,
                )
                return repr(val)

            return {
                f.name: _convert(getattr(self, f.name))
                for f in fields(self)

                    for index, item in enumerate(val):
                        item_fqn = f"{fqn}.{index}"
                        if isinstance(item, config_cls):
                            yield item_fqn, item, val, index
                        elif hasattr(item, "traverse"):
                            yield from item.traverse(config_cls, _prefix=item_fqn)
                elif hasattr(val, "traverse"):
                    yield from val.traverse(config_cls, _prefix=fqn)

        def build(self, **kwargs):
            """Construct the :class:`Configurable` subclass bound to this config.

            f"{owner_name}.Config must use @dataclass(kw_only=True, slots=True)"
        )
    for field in fields(config_cls):
        if field.init and not field.kw_only:
            raise TypeError(
                f"{owner_name}.Config field '{field.name}' must be keyword-only"
            )

        Example::

            LoRAConverter.Config(rank=16).build().convert(llama_cfg)
        """
        raise NotImplementedError


class BaseModel(Module):
    """Base class for declarative whole-model components.

        Example::

            model.init_weights(buffer_device=torch.device("npu", 0))
        """
        buffer_device = kwargs.get("buffer_device")
        self.init_states(buffer_device=buffer_device)

    def verify_module_protocol(self) -> None:
        """Ensure every submodule is a :class:`~hyper_parallel.dmodule.module.Module`.

                cfg = LlamaModel.Config(num_layers=32)
                cfg.update_from_config(trainer_config=trainer_cfg)
                model = cfg.build()
            """
            del trainer_config, kwargs

        def get_nparams_and_flops(self, model: Module, seq_len: int) -> tuple[int, int]:
            """Return parameter count and estimated FLOPs for logging.

            Example::

                nparams, nflops = cfg.get_nparams_and_flops(model, seq_len=4096)
            """
            del model, seq_len
            raise NotImplementedError(
                f"{type(self).__name__} must implement get_nparams_and_flops"
            )

_created_classes: dict[type, type] = {}


def _get_attr_by_path(obj: Any, path: str) -> Any:
    parts = path.split(".")
    for part in parts[:-1]:
        obj = getattr(obj, part)
    return getattr(obj, parts[-1])


def _set_param_by_path(module: _PlatformModule, path: str, param: Any) -> None:
    parts = path.split(".")
    if len(parts) == 1:
        module.register_parameter(parts[0], param)
        return
    parent = module
    for part in parts[:-1]:
        parent = getattr(parent, part)
    parent.register_parameter(parts[-1], param)


def _placements_equal(
    left: tuple[Placement, ...] | list[Placement],

def _placements_equal(
    left: tuple[Placement, ...] | list[Placement],
    right: list[Placement],
) -> bool:
    if len(left) != len(right):
        return False
    return all(a == b for a, b in zip(left, right))


class Module(_PlatformModule, Configurable):
    """Declarative distributed layer base class.

            instance = Configurable.Config.build(self, **kwargs)
            if self.param_init is not None:
                instance._param_init = self.param_init
            if self.sharding_config is not None:
                instance._sharding_config = self.sharding_config
            return instance

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)

            child = queue.pop(0)
            if isinstance(child, Module):
                child.init_states(buffer_device=buffer_device)
            else:
                queue.extend(child.children())

        self._init_self_parameters()

        dtensor_meta = {

            if isinstance(buf, DTensor)
        }
        self._init_self_buffers(buffer_device=buffer_device)
        for name, (mesh, placements) in dtensor_meta.items():
            new_buf = self._buffers.get(name)
            if new_buf is None or isinstance(new_buf, DTensor):
                continue
            persistent = name not in self._non_persistent_buffers_set
            self.register_buffer(
                name,
                distribute_tensor(new_buf, mesh, list(placements)),
                persistent=persistent,
            )

                f"No param_init found for parameter '{name}' in "
                f"{type(self).__name__}. Set param_init on this module's Config."
            )
        if name not in self._param_init:
            raise ValueError(
                f"No initializer for parameter '{name}' in {type(self).__name__}. "
                f"Available: {list(self._param_init.keys())}"
            )
        self._param_init[name](param)

            model.init_states()
            model.parallelize(mesh)
        """
        if self._parallelized:
            raise ValueError(
                f"{type(self).__name__} has already been parallelized. "
                "Module.parallelize() must be called at most once per instance."
            )
        self._parallelized = True

        self._parallelized = True

        sc = self.sharding_config
        if sc is None:
            for child in self.children():
                if isinstance(child, Module):
                    child.parallelize(tp_mesh)
            return

        if sc.local_map is not None:
            raise NotImplementedError("local_map will be added in M9")

        mesh_axis_names = tuple(tp_mesh.mesh_dim_names or ())
        if not mesh_axis_names:
            raise ValueError("DeviceMesh must have mesh_dim_names for parallelize()")

        self._shard_states(tp_mesh, sc, mesh_axis_names)
        self._cache_pos_arg_names()
        unbound_forward = type(self).forward

        def forward_with_redistribution(*args, **kwargs):
            args, kwargs = self._redistribute_inputs(tp_mesh, mesh_axis_names, sc, args, kwargs)
            outputs = unbound_forward(self, *args, **kwargs)
            return self._redistribute_outputs(tp_mesh, mesh_axis_names, sc, outputs)

        self.forward = forward_with_redistribution  # type: ignore[method-assign]

        for child in self.children():
            if isinstance(child, Module):
                child.parallelize(tp_mesh)

    def _shard_states(
        self,
        tp_mesh: DeviceMesh,

        sharding_config: ShardingConfig,
        mesh_axis_names: tuple[str, ...],
    ) -> None:
        """Shard parameters listed in ``sharding_config.state_shardings``."""
        for path, named_placements in sharding_config.state_shardings.items():
            param = _get_attr_by_path(self, path)
            placements = resolve_placements(named_placements, mesh_axis_names)
            if isinstance(param, DTensor):
                if not _placements_equal(tuple(param.placements), placements):
                    raise ValueError(
                        f"{type(self).__name__}.{path} is already a DTensor with "
                        f"placements {param.placements}, but sharding_config expects "
                        f"{placements}."
                    )
                continue
            tensor = param.data if hasattr(param, "data") else param
            new_local = distribute_tensor(tensor, tp_mesh, placements)
            requires_grad = getattr(param, "requires_grad", True)
            _set_param_by_path(
                self,
                path,
                platform.Parameter(new_local, requires_grad=requires_grad),
            )

        args: tuple,
        kwargs: dict,
    ) -> tuple[tuple, dict]:
        """Redistribute forward inputs per ``in_src_shardings`` / ``in_dst_shardings``."""
        if (
            sharding_config.in_dst_shardings is None
            and sharding_config.in_src_shardings is None
        ):
            return args, kwargs

        pos_arg_names = [
            name for name in self._cache_pos_arg_names() if name not in kwargs
        ]
        new_kwargs = dict(zip(pos_arg_names, args))
        new_kwargs.update(kwargs)

        in_dst_shardings = sharding_config.in_dst_shardings or {}
        in_src_shardings = sharding_config.in_src_shardings or {}

        for name, value in new_kwargs.items():
            if not platform.is_tensor(value) and not isinstance(value, DTensor):
                continue
            src_named = in_src_shardings.get(name)
            dst_named = in_dst_shardings.get(name)
            if src_named is None and dst_named is None:
                continue

            if not isinstance(value, DTensor):
                if src_named is not None:
                    layout = resolve_placements(src_named, mesh_axis_names)
                    value = DTensor.from_local(value, tp_mesh, layout)
                elif dst_named is not None:
                    layout = resolve_placements(dst_named, mesh_axis_names)
                    value = DTensor.from_local(value, tp_mesh, layout)

            if dst_named is not None and isinstance(value, DTensor):
                desired = resolve_placements(dst_named, mesh_axis_names)
                if not _placements_equal(tuple(value.placements), desired):
                    value = value.redistribute(tp_mesh, desired)

            new_kwargs[name] = value

        new_args = tuple(new_kwargs.pop(name) for name in pos_arg_names)
        return new_args, new_kwargs

    def _redistribute_outputs(
        self,
        tp_mesh: DeviceMesh,

        sharding_config: ShardingConfig,
        outputs: Any,
    ) -> Any:
        """Redistribute forward outputs per ``out_dst_shardings``."""
        out_named = sharding_config.out_dst_shardings
        if out_named is None:
            return outputs
        if not isinstance(outputs, DTensor):
            return outputs
        desired = resolve_placements(out_named, mesh_axis_names)
        if not _placements_equal(tuple(outputs.placements), desired):
            outputs = outputs.redistribute(tp_mesh, desired)
        return outputs

    @classmethod
    def from_nn_module(cls, nn_module_cls: type) -> type["Module"]:
        """Wrap a platform ``nn.*`` class as a :class:`Module` subclass.

    """Return ``torch.nn.{ModuleList,ModuleDict,Sequential}`` (requires PyTorch, lazy)."""
    # pylint: disable=C0415
    try:
        import torch.nn as nn
    except ImportError as exc:
        raise NotImplementedError(
            f"{kind} container wrappers require PyTorch (torch.nn) in M1"
        ) from exc

    mapping = {

        "ModuleDict": nn.ModuleDict,
        "Sequential": nn.Sequential,
    }
    if kind not in mapping:
        raise ValueError(f"Unknown container kind: {kind}")
    return mapping[kind]


_LAZY_CONTAINER_NAMES = frozenset({"ModuleList", "ModuleDict", "Sequential"})