get_batch

gpjax.fit.get_batch(train_data, batch_size, key)

Batch the data into mini-batches. Sampling is done with replacement.

Parameters:

• train_data (Dataset) – The training dataset.

• batch_size (int) – The batch size.

• key (KeyArray) – The random key to use for the batch selection.

Return type:

Dataset

Example

>>> import gpjax as gpx
>>> import jax.numpy as jnp
>>> import jax.random as jr

>>> X = jnp.linspace(0, 1, 100).reshape(-1, 1)
>>> y = jnp.sin(X)
>>> D = gpx.Dataset(X=X, y=y)

>>> from gpjax.fit import get_batch
>>> batch = get_batch(D, batch_size=16, key=jr.key(0))

Returns:

The batched dataset.

Return type:

Dataset

Parameters:

• train_data (Dataset)

• batch_size (int)

• key (UInt32[jaxlib._jax.Array, '2'] | Key[jaxlib._jax.Array, ''])

Expand for references to gpjax.fit.get_batch

get_batch