import torch


def tangent_basis(z):
    """
    Map (m-) coordinates z (batch_size, c-1, ...) to tangent space vector v (batch_size, c, ...).
    """
    v_shape = list(z.shape)
    v_shape[1] += 1
    v = torch.zeros(*v_shape, device=z.device)
    v[:,:-1,...] += z
    v[:,-1,...] -= z.sum(dim=1)
    return v