Function broadcast

fast_transformers/aggregate/__init__.py:42–56

(Y, G, F, X=None)

def broadcast(Y, G, F, X=None):
    device = Y.device
    if X is None:
        X = torch.zeros(
            G.shape + (Y.shape[-1],),
            device=device,
            dtype=Y.dtype
        )

    if device.type == "cpu":
        broadcast_cpu(Y, G, F, X)
    else:
        broadcast_gpu(Y, G, F, X)

    return X


# Divide the cluster into groups of equal size