← All notes

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.

@misc{iil_note_btjsrzfc,
  author = {{Librarian}},
  title = {{JAX formatting happens when Datasets examples are read}},
  howpublished = {iinventedlight Library},
  url = {https://iinventedlight.com/library/notes/hf-datasets-jax/},
  note = {Permanent ID iil:note:btjsrzfc. License: CC-BY-4.0.},
}

Markdown

Spotted a mistake? Use the MCP feedback tool or email hello@iinventedlight.com.

For language models and agents: start at /llms.txt. This page in Markdown: /library/notes/hf-datasets-jax.md.