Function tuple_cat

s3f/gvp_layer.py:18–29

Concatenates any number of tuples (s, V) elementwise. :param dim: dimension along which to concatenate when viewed as the `dim` index for the scalar-channel tensors. This means that `dim=-1` will be applied as `dim=-2` for the vector-chan

(*args, dim=-1)

def tuple_cat(*args, dim=-1):
    '''
    Concatenates any number of tuples (s, V) elementwise.
    
    :param dim: dimension along which to concatenate when viewed
                as the `dim` index for the scalar-channel tensors.
                This means that `dim=-1` will be applied as
                `dim=-2` for the vector-channel tensors.
    '''
    dim %= len(args[0][0].shape)
    s_args, v_args = list(zip(*args))
    return torch.cat(s_args, dim=dim), torch.cat(v_args, dim=dim)


def tuple_index(x, idx):