Add PyTorch TensorDict validation backend - #2403
Open
cosmicBboy wants to merge 3 commits into
Open
Conversation
Squashed rebase of dev/pytorch-tensordict-phase4 (phases 1-4 of the TensorDict integration) onto main, reviewed and updated against tensordict 0.13 / torch 2.12. Integration surface: - TensorDictSchema / Tensor components with dtype, shape, batch_size, and value-check validation for TensorDict and tensorclass objects - Class-based TensorDictModel API with Field annotations and Config - dtype coercion (schema-level and per-tensor) - Schema inference from TensorDict/tensorclass data - YAML/JSON schema serialization plus save/load with embedded schema - Hypothesis strategies for TensorDict and tensorclass generation - Docs guide (docs/source/pytorch_guide) and API reference Updates on top of the original branch work: - Use tensordict's canonical type predicates (is_tensor_collection, is_tensorclass) and register the backend on TensorDictBase so all subclasses validate, including LazyStackedTensorDict - Fix save() on plain dicts (auto_batch_size_ inference) and load()'s broken infer_schema reference - Fix TensorDictModel Config options (coerce, name, title, description) and per-field nullable/coerce/required being silently dropped - Honor per-tensor coerce=True when schema-level coerce is off - Rewrite strategies on hypothesis.extra.numpy: exact dtype fidelity, batch-dim alignment, and check-aware element generation (isin, ge/gt/le/lt, in_range) - Register modern torch dtypes (uint16/32/64, float8) and drop deprecated quantized dtypes; fix engine handling of unknown dtypes - Keep pandera.io/pandera.strategies/pandera.schema_inference lazy so importing pandera does not require torch - Remove committed merge-conflict markers from docs and fix doc examples (all 41 code cells in the pytorch guide now execute) - Regression tests for LazyStackedTensorDict validation, dict save round-trip, and strategy dtype/batch fidelity (90 tests total) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: cosmicBboy <niels.bantilan@gmail.com>
pyspark 4.2.0 (released 2026-07-14) leaks "Worker Monitor for python3" threads for reused python workers — a regression of SPARK-35009. On macOS CI runners the Spark JVM exhausts the per-process thread limit mid-suite and dies with "OutOfMemoryError: unable to create native thread"; every subsequent test in tests/pyspark/ test_schemas_on_pyspark_pandas.py then fails with ConnectionRefusedError. Ubuntu runners have higher thread limits and pass. The same PR head passed CI on pyspark 4.1.2 before 4.2.0 came out, so this is dependency drift, not a code issue. Pin pyspark to < 4.2 in the nox testing requirements (CI-only; the user-facing pyspark extra is unchanged since pandera itself works with 4.2) until the thread leak is fixed upstream. Also copy the pyproject dependencies list in _testing_requirements before appending extras — the previous in-place += mutated the shared PYPROJECT dict, leaking one session's requirements into the next when multiple nox sessions run in one process. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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.
Summary
Brings the TensorDict integration work from
dev/pytorch-tensordict-phase4(phases 1–4) onto latestmainas a single squashed change, reviewed and updated against the current state of the pytorch/tensordict library (tensordict 0.13, torch 2.12).Integration surface
TensorDictSchema/Tensorcomponents validating dtype, shape,batch_size, and value checks forTensorDictandtensorclassobjectsTensorDictModelAPI withFieldannotations andConfigpa.infer_schema)save/loadwith embedded schema metadatadocs/source/pytorch_guide/) and API reference pagetorchextra:pip install 'pandera[torch]'Updates on top of the original branch
Reviewed against tensordict 0.13 and fixed issues found along the way:
is_tensor_collection/is_tensorclasspredicates and register the validation backend onTensorDictBase, so all subclasses validate — includingLazyStackedTensorDict(TensorDict.lazy_stackresults are notisinstanceofTensorDict, so the original code missed them)save()on a plain dict previously crashed (bogusEngine.dtype("tensordict")call) — now builds aTensorDictand infers batch dims viaauto_batch_size_();load()referenced an unimportedinfer_schemaConfig.coerce(andname/title/description) and per-fieldnullable/coerce/requiredwere silently dropped when building the schema; per-tensorcoerce=Truenow works without schema-level coercehypothesis.extra.numpywith exact dtype fidelity (was hardcoded float32/int64), batch-dim alignment (was generating tensors inconsistent withbatch_size, a hardRuntimeErrorin current tensordict), and check-aware element generation (isin,ge/gt/le/lt,in_range)uint16/32/64,float8_e4m3fn,float8_e5m2), drop deprecated quantized dtypes, and raise a clear error for unknown dtype inputspandera.io,pandera.strategies, andpandera.schema_inferencestay lazy — importing pandera does not import torch (strategies exports use module__getattr__)tensordict_io.md, and fixed several doc examples that were broken or numerically unsound; all 41 executable code cells in the pytorch guide now run cleanly against tensordict 0.13Test plan
tests/tensordict/— 90 tests pass (84 from the original branch + 6 new regression tests for LazyStackedTensorDict validation, plain-dictsave()round-trip, and strategy dtype/batch fidelity)docs/source/pytorch_guide/execute cleanly (3 consecutive runs to shake out hypothesis flakiness)tests/pandas/test_schemas.pypasses with torch installed (no cross-backend regressions)import panderaremains torch-free🤖 Generated with Claude Code