import torch


def _r(A, b_0, B, y):
    """Compute y - b_0 - sum_j A_j@b_j."""
    # bmm will return a tensor with individual A_j@b_j products.
    AB = torch.bmm(A.permute(0, 2, 1), B.unsqueeze(2)).sum(dim=0).squeeze()
    return y - b_0 - AB