def materialize_tensor(
    tensor: torch.Tensor,
) -> Generator[torch.Tensor, None, None]:
    """Context manager for initializing DTensor parameters.

    This context manager temporarily materializes a DTensor as a full tensor for initialization,
    then redistributes the initialized values back to the original DTensor.

    NOTE: This approach is not memory-efficient as it materializes a full tensor on every rank
    and then redistributes it back to a DTensor. However, since model initialization happens
    only once, this overhead is generally acceptable.

    Args:
        tensor: The tensor to materialize. Can be either a regular Tensor or a DTensor.

    Yields:
        torch.Tensor: The materialized full tensor that can be initialized.
    """
    is_dtensor = isinstance(tensor, DTensor)

    with torch.no_grad():
        if is_dtensor:
            full_tensor = tensor.full_tensor()
        else:
            full_tensor = tensor

        yield full_tensor
        if is_dtensor:
            # Redistribute the updated full tensor back to the original DTensor
            with torch.no_grad():
                temp_tensor = distribute_tensor(
                    full_tensor,
                    device_mesh=tensor.device_mesh,
                    placements=tensor.placements,
                )
                tensor.to_local().copy_(temp_tensor.to_local())
Context manager for initializing DTensor parameters.

    This context manager temporarily materializes a DTensor as a full tensor for initialization,
    then redistributes the initialized values back to the original DTensor.

    NOTE: This approach is not memory-efficient as it materializes a full tensor on every rank
    and then redistributes it back to a DTensor. However, since model initialization happens
    only once, this overhead is generally acceptable.

    Args:
        tensor: The tensor to materialize. Can be either a regular Tensor or a DTensor.

    Yields:
        torch.Tensor: The materialized full tensor that can be initialized.
    
llmfoundry/models/utils/param_init_fns.py