//  [python]
def Batch_OMP(data, dictionary, max_nonzero, tolerance=1e-7, debug=False):
    """
    for details on variable names, see
    https://sparse-plex.readthedocs.io/en/latest/book/pursuit/omp/batch_omp.html
    or the original paper
    http://www.cs.technion.ac.il/~ronrubin/Publications/KSVD-OMP-v2.pdf

    NOTE - the implementation below works on transposed versions of the input signal to make the batch size the first
           coordinate, which is how pytorch expects the data..
    """
    vector_dim, batch_size = data.size()
    dictionary_t = dictionary.t()
    G = dictionary_t.mm(dictionary)  # this is a Gram matrix
    eps = torch.linalg.norm(data, dim=0)  # the residual, initalized as the L2 norm of the signal
    h_bar = dictionary_t.mm(data).t()  # initial correlation vector, transposed to make batch_size the first dimension

    # note - below this line we no longer use "data" or "dictionary"

    h = h_bar
    x = torch.zeros_like(h_bar)  # the resulting sparse code
    L = torch.ones(batch_size, 1, 1, device=h.device)  # Contains the progressive Cholesky of G in selected indices
    I = torch.ones(batch_size, 0, device=h.device).long()
    I_logic = torch.zeros_like(h_bar).bool()  # used to zero our elements is h before argmax
    delta = torch.zeros(batch_size, device=h.device)  # to track errors

    k = 0
    while k < max_nonzero and eps.max() > tolerance:
        k += 1
        # use "I_logic" to make sure we do not select same index twice
        index = (h*(~I_logic).float()).abs().argmax(dim=1)  # todo - can we use "I" rather than "I_logic"
        _update_logical(I_logic, index)
        batch_idx = torch.arange(batch_size, device=G.device)
        expanded_batch_idx = batch_idx.unsqueeze(0).expand(k, batch_size).t()

        if k > 1:  # Cholesky update
            # Following line is equivalent to:
            #   G_stack = torch.stack([G[I[i, :], index[i]] for i in range(batch_size)], dim=0).view(batch_size, k-1, 1)
            G_stack = G[I[batch_idx, :], index[expanded_batch_idx[...,:-1]]].view(batch_size, k-1, 1)
            w = torch.triangular_solve(G_stack, L, upper=False, ).solution.view(-1, 1, k-1)
            w_corner = torch.sqrt(1-(w**2).sum(dim=2, keepdim=True))  # <- L corner element: sqrt(1- w.t().mm(w))

            # do concatenation into the new Cholesky: L <- [[L, 0], [w, w_corner]]
            k_zeros = torch.zeros(batch_size, k-1, 1, device=h.device)
            L = torch.cat((
                torch.cat((L, k_zeros), dim=2),
                torch.cat((w, w_corner), dim=2),
            ), dim=1)

        # update non-zero indices
        I = torch.cat([I, index.unsqueeze(1)], dim=1)

        # x = solve L
        # The following line is equivalent to:
        #   h_stack = torch.stack([h_bar[i, I[i, :]] for i in range(batch_size)]).unsqueeze(2)
        h_stack = h_bar[expanded_batch_idx, I[batch_idx, :]].view(batch_size, k, 1)
        x_stack = torch.cholesky_solve(h_stack, L)

        # de-stack x into the non-zero elements
        # The following line is equivalent to:
        #   for i in range(batch_size):
        #       x[i:i+1, I[i, :]] = x_stack[i, :].t()
        x[batch_idx.unsqueeze(1), I[batch_idx]] = x_stack[batch_idx].squeeze(-1)

        # beta = G_I * x_I
        # The following line is equivalent to:
        # beta = torch.cat([x[i:i+1, I[i, :]].mm(G[I[i, :], :]) for i in range(batch_size)], dim=0)
        beta = x[batch_idx.unsqueeze(1), I[batch_idx]].unsqueeze(1).bmm(G[I[batch_idx], :]).squeeze(1)

        h = h_bar - beta

        # update residual
        new_delta = (x * beta).sum(dim=1)
        eps += delta-new_delta
        delta = new_delta

        if debug and k % 1 == 0:
            print('Step {}, residual: {:.4f}, below tolerance: {:.4f}'.format(k, eps.max(), (eps < tolerance).float().mean().item()))

    return x.t()