Enable usage of the pipeline with JAX models

[bnubald]

To enable models built using JAX to use the pipeline for end to end runs.

Relates to icenet-ai/icenet-roadmap#57

Due to JAX focusing on the model development side of things, could piggyback on a lot of the existing work done by @JimCircadian for data loading since JAX supports it.

Refer to:
https://jax.readthedocs.io/en/latest/notebooks/neural_network_with_tfds_data.html#data-loading-with-tensorflow-datasets