import torch


def _compute_dx(x, grid):
    """
        x: [n_batch, n_dim]
        grid: [n_batch, n_dim, n_grid]
        
        returns
            dx: [n_batch, n_dim, n_grid]
    """

    x = x.unsqueeze(-1)
    #     grid = grid.unsqueeze(0)

    dx = x - grid
    dx = torch.cat([dx[..., [0]], dx], dim=-1)

    return dx