def _init_split_assignments(
        self,
        model: Union[HfModelHandler, PyTorchModelHandler],
        pytorch_model: "nn.Module",
        unique_embeds_lm_head_splits: bool,
    ) -> tuple[dict[str, int], int, set]:
        """Initialize the split assignments for the model.

        :param model: The input model to split.
        :param pytorch_model: The loaded input model.
        :param unique_embeds_lm_head_splits: Whether to assign embedding and lm_head layers to their own splits.
        :return: A dictionary of split assignments, next split index, and set of modules to exclude.
        """
        split_assignments = {}
        split_idx = 0
        modules_to_exclude = set()
        if unique_embeds_lm_head_splits:
            # assign embedding layer to its own split
            if isinstance(model, HfModelHandler):
                embed_names = ModelWrapper.from_model(pytorch_model).get_embeds()[1]
            else:
                embed_names = ["model.embed_tokens"]
            for embed_name in embed_names:
                split_assignments[embed_name] = split_idx
            split_idx += 1

            # exclude embedding and lm head from split calculation
            modules_to_exclude.update([*embed_names, "lm_head"])

        return split_assignments, split_idx, modules_to_exclude
Initialize the split assignments for the model.

        :param model: The input model to split.
        :param pytorch_model: The loaded input model.
        :param unique_embeds_lm_head_splits: Whether to assign embedding and lm_head layers to their own splits.
        :return: A dictionary of split assignments, next split index, and set of modules to exclude.
        
olive/passes/pytorch/capture_split_info.py