Function batched_index_select

utils/general_util.py:19–22

(t, dim, inds)

    return output_start + ((output_end - output_start) / (input_end - input_start)) * (input_clamped - input_start)

def batched_index_select(t, dim, inds):
    dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2))
    out = t.gather(dim, dummy) # b x e x f
    return out

def tridiagonal_solve(b, A_upper, A_diagonal, A_lower):
    """Solves a tridiagonal system Ax = b.