import torch


def uniform_attention(q, v):
    """
    Uniform attention. Equivalent to Neural Process.

    Parameters
    ----------
    q : torch.Tensor
        Shape (batch_size, m, k_dim)

    v : torch.Tensor
        Shape (batch_size, n, v_dim)

    Returns
    -------
    r : torch.Tensor
        Shape (batch_size, m, v_dim)
    """
    r = torch.mean(v, dim=1, keepdim=True) # shape [B, 1, v_dim]
    # Infer the number of target points from query
    m = q.shape[1]
    # Reshape the output
    r = torch.repeat_interleave(r, m, dim=1)
    return r