I added tests and a minimal example notebook.

Tests are passing on torch and tensorflow, but fail on jax.

@LarsKue since you are the architect of the stateless_compute_metrics, could you look into how we can make this work for jax?

The final section of the keras guide on custom training loops in jax proves that this can be rather straight forward, but I am unsure how to implement it in our case: https://keras.io/guides/writing_a_custom_training_loop_in_jax/