Add a Tensor feature for arrays with fixed or dynamic shapes - #8603
Open
behroozazarkhalili wants to merge 5 commits into
Open
behroozazarkhalili wants to merge 5 commits into
behroozazarkhalili wants to merge 5 commits into
Conversation
Array2D..Array5D require a fixed shape declared in the schema, which rules out data such as multi-channel telescope frames whose dimensions vary from row to row. Tensor(shape, dtype) stores a fully specified shape as Arrow's canonical arrow.fixed_shape_tensor extension and any shape with unknown dimensions (or an explicit ndim) as the canonical arrow.variable_shape_tensor layout, struct<data: list<T>, shape: fixed_size_list<int32>[ndim]> with the spec's metadata, so files written by datasets read back in plain pyarrow as the canonical types. Encoding accepts numpy arrays and nested lists, validates dtype and every fixed dimension, keeps integer precision (int64 and uint64 values never pass through float64) and rejects out-of-range or non-boolean values. Decoding returns numpy arrays; the numpy, torch, tensorflow, jax, pandas and polars formatters return native tensors with the declared dtype, including for tensors nested in List and dict columns and for None rows. Tensor takes part in the standard encoding pipeline, so sibling fields, new keys added by map, Json siblings, LargeList nesting, torch tensors that require grad, and the arrow writer's type inference behave as they do for other features; the extension survives embed_table_storage. Features.from_arrow_schema recognises the canonical types whether or not datasets was imported before the schema was read. Array2D..Array5D are unchanged. Content-defined-chunking friendly layouts were not measured; the canonical layout was verified by reading the written IPC and Parquet files in a pyarrow-only process. Resolves huggingface#7738
added 3 commits
September 11, 2026 05:46
…ypes On pyarrow builds where arrow.fixed_shape_tensor and arrow.variable_shape_tensor are registered by Arrow C++ (as on CI), reading a schema yields the native extension class instead of the datasets Tensor class. Casting between the two failed with an extension-to-extension cast error. Recognise the canonical types by name and storage, rebuild the Tensor type from their metadata, and rewrap storage instead of casting across extension classes. The subprocess interop test no longer requires datasets to be absent from the pyarrow site-packages, which is the CI layout.
pyarrow 24, the minimum supported version, rejects Parquet writes of
list arrays whose null slots span values ("Lists with non-zero length
null components are not supported"). The storage rewrapping introduced
for natively registered canonical types produced such slots for null
rows. Build null rows with empty spans, matching what Arrow's own
canonical tensor types emit, and cover the plain-pyarrow reader path.
The newest jax removed the jax.experimental.enable_x64 context manager used by the dtype-preservation tests, which broke the deps-latest CI lanes. Flip the flag with jax.config.update and restore the previous value in a finally block, which works on every supported jax version.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resolves #7738.
This adds a
Tensor(shape, dtype)feature for multi-dimensional arrays whose dimensions can vary from row to row, following the direction discussed in the issue: fixed shapes use Arrow's canonicalarrow.fixed_shape_tensorextension, and shapes with unknown dimensions (or an explicitndim) use the canonicalarrow.variable_shape_tensorlayout (struct<data: list<T>, shape: fixed_size_list<int32>[ndim]>with the spec's metadata). pyarrow 25 has no Python wrapper for the variable-shape type yet, so datasets registers the extension itself; files written this way read back in a pyarrow-only process as the canonical types, which is how interoperability was verified.What it does:
List(Tensor)and dict columns, withNonerows preserved.map,Jsonsiblings,LargeListnesting, torch tensors withrequires_grad, and the writer's type inference behave as for other features; the extension also survivesembed_table_storage.Features.from_arrow_schemarecognises the canonical types regardless of import order.Array2D..Array5Dare unchanged.Tests live in
tests/features/test_tensor.py(fixed and dynamic shapes, dtype validation, nulls, batched paths, save_to_disk / parquet round trips, each formatter, precision, nested layouts, canonical interop from a pyarrow-only subprocess).