code: def dataset(batch_size):
      while True:
          idx = np.random.choice(len(X), size=batch_size)
          yield X[idx], Y[idx]
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/02_lifted_transforms.py
imports: import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling a global function with example usage: dataset(batch_size) without return types