fix: be explicit about squeeze dim in prioritised sampling to avoid flattening (1,1) arrays

[callumtilbury]

Currently, the following snippet will fail:

from flashbax import make_prioritised_flat_buffer
import jax
import jax.numpy as jnp

buffer = make_prioritised_flat_buffer(
    max_length=100,
    min_length=1,
    sample_batch_size=1,  # NB
    add_sequences=False,
)

timestep = {"obs": jnp.zeros(shape=(3)),}

state = buffer.init(timestep)

for i in range(5):
    timestep = {
        "obs": jnp.ones(shape=(3)) * i,
    }
    state = buffer.add(state, timestep)

buffer.sample(state, jax.random.PRNGKey(0))

because of this line:

flashbax/flashbax/buffers/sum_tree.py

Line 211 in b556e15

state, None, query_values.squeeze()

If the sample_batch_size is 1, query_values is shape (1,1), which squeezes to ().

Instead we must be explicit about the squeeze dim.