def summon_dtensor(init_fn: Callable) -> Callable:
    """Unshard and Reshard DTensor parameters for initialization.

    This decorator wraps an initialization function to handle both regular tensors/modules
    and those containing DTensor parameters by temporarily materializing DTensors as full
    tensors during initialization.

    Args:
        init_fn: The initialization function to wrap.

    Returns:
        A wrapped initialization function that can handle DTensor parameters.
    """

    def init_fn_wrapper(
        obj: nn.Module | torch.Tensor,
        *args: Any,
        **kwargs: Any,
    ) -> bool:
        if isinstance(obj, nn.Module):
            with materialize_module(obj) as obj:
                return init_fn(obj, *args, **kwargs)
        elif isinstance(obj, torch.Tensor):
            with materialize_tensor(obj) as obj:
                return init_fn(obj, *args, **kwargs)
        else:
            raise TypeError(f'Invalid object type: {type(obj)}')

    return init_fn_wrapper
Unshard and Reshard DTensor parameters for initialization.

    This decorator wraps an initialization function to handle both regular tensors/modules
    and those containing DTensor parameters by temporarily materializing DTensors as full
    tensors during initialization.

    Args:
        init_fn: The initialization function to wrap.

    Returns:
        A wrapped initialization function that can handle DTensor parameters.
    
llmfoundry/models/utils/param_init_fns.py