utils.py

__all__ = ['create_batches']

from typing import Generator, Union
import numpy as np              # Array construction
from numpy.typing import NDArray
import jax                      # JAX for numpy
import jax.numpy as jnp         # JAX for numpy
import equinox as eqx           # Equinox for JAX
from jax import lax             # JAX for linear algebra

@eqx.filter_jit
def create_batches(x_data, y_data, batch_size, key):
    """Create mini-batches from training data.
    
    Shuffles data and splits into fixed-size batches using JAX scan.
    
    Args:
        x_data: Input data array
        y_data: Target data array
        batch_size: Size of each mini-batch
        key: JAX random key for shuffling
    
    Returns:
        Tuple of (x_batches, y_batches) where each is an array of batches
        
    Note:
        Author: Atharva Hans
    """
    indices = jnp.arange(len(y_data))
    shuffled_indices = jax.random.permutation(key, indices, independent=True)
    num_batches = len(y_data) // batch_size
    def step_scan(carry, _):
        iter_ = carry
        batch_indices = lax.dynamic_slice(
            shuffled_indices, (iter_ * batch_size,), (batch_size,)
        )
        batch_x = x_data[batch_indices]
        batch_y = y_data[batch_indices]
        return iter_ + 1, (batch_x, batch_y)
    init = 0
    _, batches = lax.scan(step_scan, init, None, length=int(num_batches))
    return batches