Running JAX in DALI pipeline

See the plugin.jax.fn.jax_function for documentation and examples on running custom JAX augmentations as a part of DALI iterator or pipeline.