import torch


def _make_aligned_labels(inputs):
    """Uses the shape of inputs to infer batch_size, n_classes, and n_sample_per_class. From this, we build the one-hot
    encoding label tensor aligned with inputs. This is used to keep the lable information when tensors are flatenned
    across the n_class and n_sample_per_class.

    Args:
        inputs: tensor of shape (batch, n_classes, n_sample_per_class, z_dim) containing encoded examples for a task.
    Returns:
        tensor of shape (batch, n_classes, n_sample_pc, n_classes) containing the one-hot encoding label of each example
    """
    batch, n_classes, n_sample_pc, z_dim = inputs.shape
    identity = torch.eye(n_classes, dtype=inputs.dtype, device=inputs.device)
    return identity[None, :, None, :].expand(batch, -1, n_sample_pc, -1).contiguous()