Applies Transformer model on the inputs. Args: x: Inputs to the layer. train: Set to `True` when training. Retur

Applies Transformer MlpBlock module.

Embed only, or complete forward pass. Args: tokens: Embedded, then and appended to `embedded_prefix`. Can be None. embedd

Create a PyTorch data loader. Args: dataset: The dataset to load. local_batch_size: The local batch size for each pro

Collate the batch elements into batched numpy arrays.

Tell JAX inside the worker process not to preallocate the GPU memory.

Converts a PyTree of arrays into a human-readable string for logging.

Get the assets directory for this config.

Checks that two PyTrees have the same structure and optionally checks shapes and dtypes. Creates a much nicer error message than if `jax.tree.map`

Get the checkpoint directory for this config.

Compose a sequence of transforms into a single transform.

Converts an image to uint8 if it is a float image. This is important for reducing the size of the image when sending it over the network.

Create a data config.

Create a new model, initializing parameters.