keras/src/backend/jax/numpy.py/squeeze
def squeeze(x, axis=None):
    if isinstance(x, jax_sparse.BCOO):
        if axis is None:
            axis = tuple(i for i, d in enumerate(x.shape) if d == 1)
        axis = to_tuple_or_list(axis)
        return jax_sparse.bcoo_squeeze(x, dimensions=axis)
    return jnp.squeeze(x, axis=axis)