feat: simple utility to get timestep count

[callumtilbury]

From discussions in #33.

Demo:

import flashbax as fbx
import jax.numpy as jnp
from flashbax.utils import get_timestep_count

buffer = fbx.make_flat_buffer(
    max_length=10,
    min_length=1,
    sample_batch_size=1,
    add_batch_size=2,
    add_sequences=False,
)

timestep = {"obs": jnp.array([1.0])}

state = buffer.init(timestep)

for i in range(10):
    timestep = {
        "obs": jnp.tile(jnp.array([i], dtype=jnp.float32), (2, 1)),
    }
    state = buffer.add(state, timestep)

    print("Timestep count =", get_timestep_count(state))

Output:

Timestep count = 2
Timestep count = 4
Timestep count = 6
Timestep count = 8
Timestep count = 10
Timestep count = 10
Timestep count = 10
Timestep count = 10
Timestep count = 10
Timestep count = 10

[EdanToledo]

@callumtilbury I'm happy with this but can we have a small test that it is jittable - i dont even really need the test to be in the codebase, i just want confirmation.

[callumtilbury]

@EdanToledo indeed, previously confirmed jittable!

#33 (comment)

[EdanToledo]

then im happy

[EdanToledo]

LGTM!