import torch


def _compute_bin_filter(x, grid):
    """
        x: [n_batch, n_dim]
        grid: [n_batch, n_dim, n_grid]
        
        returns
            bin_filter: [n_batch, n_dim, n_grid + 2]
    """

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

    is_in = (x >= grid[..., :-1]) & (x < grid[..., 1:])
    is_left = x < grid[..., [0]]
    is_right = x >= grid[..., [-1]]

    bin_filter = torch.cat([is_left, is_in, is_right], dim=-1).to(x)

    return bin_filter