import torch


def adaptive_add_fusion(z: torch.Tensor, p: torch.Tensor):
    """Adds a point cloud to a noise tensor that can have a smaller dimension.

    The noise tensor `z` with `dim_z` channels will be added to the first `dim_z`
    channels of the point cloud `p`.

    Args:
        z: Reference point cloud of shape [batch_size, dim_z, num_point].
        p: Reference point cloud of shape [batch_size, dim_p, num_point].

    Returns:
        Point cloud tensor of shape [batch_size, dim_p, num_point].
    """
    dim_z = z.shape[1]
    dim_p = p.shape[1]
    if dim_z <= dim_p:
        return torch.cat([z + p[:, :dim_z, :], p[:, dim_z:, :]], dim=1)
    else:
        raise RuntimeError(
            "dim_z must be less or equal to dim_p. "
            + f"got `dim_z` = {dim_z} > {dim_p} = `dim_p`",
        )