def _image_preprocess_fn(features):
      if "image" in features:
        features["image"] = tf.cast(features["image"], tf.float32) / 255.0
      if "id" in features:  # Included in some TFDS datasets, breaks JAX.
        del features["id"]
      return features