How do you create a Preprocessor using JAX instead of tensorflow?

seqio advertises its compatibility with JAX and PyTorch in the [README](https://github.com/google/seqio). However, the only [docs](https://seqio.readthedocs.io/en/latest/overview.html#preprocessors) I have found use `tf.Tensor`s as inputs and outputs for Pre- and Postprocessors. How does one create a Preprocessor that operates on `jax.Array`s, `np.ndarray`s, or anything else?