---
title: JAX formatting happens when Datasets examples are read
slug: hf-datasets-jax
license: CC-BY-4.0
license_url: https://creativecommons.org/licenses/by/4.0/
source: https://iinventedlight.com/library/notes/hf-datasets-jax.md
source_license: Apache-2.0
attribution: Original synthesis by iinventedlight, drafted by Librarian with AI assistance; based on the attributed sources
  below.
author: Librarian
client: Librarian
date: '2026-11-05'
checked_on: '2026-11-05'
as_of: '2026-11-05'
status: published
confidence: medium
category: coding-references
tags:
- hf
- datasets
- jax
untested: Reference note. No code or commands were run to validate behavior.
sources:
- id: S1
  title: Use with JAX
  publisher: Hugging Face contributors
  url: https://raw.githubusercontent.com/huggingface/datasets/main/docs/source/use_with_jax.mdx
  accessed: '2026-11-05'
  license: Apache-2.0
  license_url: https://raw.githubusercontent.com/huggingface/datasets/main/LICENSE
  license_verified: '2026-11-05'
  green: true
  attribution: Hugging Face contributors; Use with JAX; Apache-2.0
  snapshot: _fetched/hf-jax.txt
- id: S2
  title: Process
  publisher: Hugging Face contributors
  url: https://raw.githubusercontent.com/huggingface/datasets/main/docs/source/process.mdx
  accessed: '2026-11-05'
  license: Apache-2.0
  license_url: https://raw.githubusercontent.com/huggingface/datasets/main/LICENSE
  license_verified: '2026-11-05'
  green: true
  attribution: Hugging Face contributors; Process; Apache-2.0
  snapshot: _fetched/hf-process.txt
---
# JAX formatting happens when Datasets examples are read

## Summary
`with_format("jax")` configures dataset reads to return JAX arrays for supported values. The format is applied when data is accessed, rather than when the formatting call is made. [S1, S2]

## Details
Strings and binary values stay unchanged. A `DatasetDict` can apply the JAX format across its member datasets. The general formatting interface also accepts `columns=` to limit which columns it formats. `set_format` changes the existing object, whereas the guide uses `with_format` with reassignment. These are formatting choices, not a request to run model training. [S1, S2]

## Claims
- id: C1
  claim: "JAX formatting is delayed until data access."
  source: S1
  evidence: "formatting is not applied\nuntil you actually access the data"
- id: C2
  claim: "JAX formatting leaves string and binary values unchanged."
  source: S1
  evidence: "string and binary objects are unchanged"
- id: C3
  claim: "columns= limits which columns the formatter handles."
  source: S2
  evidence: "choose which the columns you want to format using `columns=`"

## Sources
- [S1] Use with JAX. Hugging Face contributors. https://raw.githubusercontent.com/huggingface/datasets/main/docs/source/use_with_jax.mdx. Apache-2.0; license: https://raw.githubusercontent.com/huggingface/datasets/main/LICENSE. Accessed 2026-11-05.
- [S2] Process. Hugging Face contributors. https://raw.githubusercontent.com/huggingface/datasets/main/docs/source/process.mdx. Apache-2.0; license: https://raw.githubusercontent.com/huggingface/datasets/main/LICENSE. Accessed 2026-11-05.

## Attribution and notices
This note is CC BY 4.0. Source material has been condensed and restated; short evidence excerpts retain their source license.

Hugging Face Datasets documentation is Apache-2.0. This note restates the documentation and includes only short evidence excerpts. See the accompanying source license.