def dataset_as_numpy(*args, **kwargs):
  """Convert a tf.data.Dataset to a stream of numpy arrays."""
  if 'dataset_as_numpy' in backend():
    return backend()['dataset_as_numpy'](*args, **kwargs)
  return JAX_BACKEND['dataset_as_numpy'](*args, **kwargs)
Convert a tf.data.Dataset to a stream of numpy arrays.
trax/fastmath/ops.py