diff --git a/.github/workflows/ci-tests.yml b/.github/workflows/ci-tests.yml index 3a882e3b6..1c1cca8d7 100644 --- a/.github/workflows/ci-tests.yml +++ b/.github/workflows/ci-tests.yml @@ -14,6 +14,7 @@ on: - dev/xarray - bugfix - "release/*" + - "dev/*" env: DEFAULT_PYTHON: 3.11 diff --git a/asv_bench/benchmarks/dataframe_schema.py b/asv_bench/benchmarks/dataframe_schema.py index 3827f10d4..21f7d2f24 100644 --- a/asv_bench/benchmarks/dataframe_schema.py +++ b/asv_bench/benchmarks/dataframe_schema.py @@ -1,21 +1,9 @@ # Airspeed Velocity Benchmarks for pandera import pandas as pd -from pandera.pandas import ( - Column, - DataFrameSchema, - Bool, - Category, - Check, - DateTime, - Float, - Int, - Object, - String, - Timedelta, - check_input, - check_output, -) +from pandera.pandas import (Bool, Category, Check, Column, DataFrameSchema, + DateTime, Float, Int, Object, String, Timedelta, + check_input, check_output) class Validate: diff --git a/asv_bench/benchmarks/series_schema.py b/asv_bench/benchmarks/series_schema.py index 30c9a1cae..cd96d31fd 100644 --- a/asv_bench/benchmarks/series_schema.py +++ b/asv_bench/benchmarks/series_schema.py @@ -1,21 +1,9 @@ # Airspeed Velocity Benchmarks for pandera import pandas as pd -from pandera.pandas import ( - Column, - DataFrameSchema, - SeriesSchema, - Bool, - Category, - Check, - DateTime, - Float, - Int, - Object, - String, - Timedelta, - String, -) +from pandera.pandas import (Bool, Category, Check, Column, DataFrameSchema, + DateTime, Float, Int, Object, SeriesSchema, String, + Timedelta) class Validate: diff --git a/docs/source/ibis.md b/docs/source/ibis.md index e69fdfbd4..edf37d179 100644 --- a/docs/source/ibis.md +++ b/docs/source/ibis.md @@ -204,7 +204,11 @@ invalid_t = ibis.memtable({ "b": ["d", "e", "f"], "c": [0.0, 1.1, -0.1], }) -ModelWithChecks.validate(invalid_t, lazy=True) + +try: + ModelWithChecks.validate(invalid_t, lazy=True) +except pa.errors.SchemaErrors as exc: + print(exc) ``` (supported-ibis-dtypes)= diff --git a/docs/source/index.md b/docs/source/index.md index 201a3cb30..42a5dac01 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -464,6 +464,7 @@ configuration supported_libraries xarray_guide/index +pytorch_guide/index integrations ``` diff --git a/docs/source/pytorch_guide/error_reporting.md b/docs/source/pytorch_guide/error_reporting.md new file mode 100644 index 000000000..9d5fac08e --- /dev/null +++ b/docs/source/pytorch_guide/error_reporting.md @@ -0,0 +1,64 @@ +--- +file_format: mystnb +--- + +(pytorch-error-reporting)= + +# Error Reporting + +Pandera raises {class}`~pandera.errors.SchemaError` for individual validation failures +and {class}`~pandera.errors.SchemaErrors` when `lazy=True`. + +## Non-lazy validation (fail-fast) + +By default, validation stops on the first error: + +```{code-cell} python +import torch +from tensordict import TensorDict +import pandera.tensordict as pa + +schema = pa.TensorDictSchema( + keys={"x": pa.Tensor(dtype=torch.float32, shape=(None, 10))}, + batch_size=(32,), +) + +td = TensorDict({"x": torch.randn(16, 10)}, batch_size=[16]) + +try: + schema.validate(td) +except pa.errors.SchemaError as exc: + print(f"Error: {exc}") + print(f"Reason: {exc.reason_code}") +``` + +## Lazy validation + +Set `lazy=True` to collect all errors: + +```{code-cell} python +td_invalid = TensorDict( + {"x": torch.randn(16, 10), "y": torch.randn(16, 10)}, + batch_size=[16], +) + +try: + schema.validate(td_invalid, lazy=True) +except pa.errors.SchemaErrors as exc: + print(f"Found {len(exc.schema_errors)} errors") + for err in exc.schema_errors: + print(f" - {str(err)}") + print(f"\nError counts: {exc.error_counts}") +``` + +## Error reason codes + +- `WRONG_DATATYPE` — tensor dtype mismatch +- `COLUMN_NOT_IN_DATAFRAME` — missing key +- `CHECK_ERROR` — value check failure + +## See also + +- {ref}`checks` — Check API +- {ref}`pytorch-tensordict-schema` — TensorDictSchema +- {ref}`lazy-validation` — general lazy validation guide diff --git a/docs/source/pytorch_guide/index.md b/docs/source/pytorch_guide/index.md new file mode 100644 index 000000000..d16c1e093 --- /dev/null +++ b/docs/source/pytorch_guide/index.md @@ -0,0 +1,94 @@ +--- +file_format: mystnb +--- + +(pytorch-guide)= + +# PyTorch Tensor Data Validation + +[PyTorch](https://pytorch.org/) provides {class}`~torch.Tensor` objects and +[{class}`~tensordict.TensorDict` for managing collections of tensors], +and [{class}`~tensordict.tensorclass`][tensorclass] for typed tensor +collections. + +Pandera validates them with the same patterns as the other dataframe backends: +schema objects, optional {class}`~pandera.api.checks.Check` instances, and +global {ref}`configuration `. + +## Installation + +```bash +pip install 'pandera[torch]' +pip install tensordict +``` + +## Quick start + +```{code-cell} python +import torch +from tensordict import TensorDict +import pandera.tensordict as pa + +schema = pa.TensorDictSchema( + keys={ + "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)), + "action": pa.Tensor(dtype=torch.float32, shape=(None, 5)), + }, + batch_size=(32,), +) + +td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + batch_size=[32], +) +schema.validate(td) +``` + +## Define a schema with a class-based model + +```{code-cell} python +class RL(pa.TensorDictModel): + """Schema for reinforcement learning data.""" + + # Use PyTorch dtypes in type annotations + observation: torch.float32 = pa.Field(shape=(None, 10)) + action: torch.int64 = pa.Field(shape=(None,)) + reward: torch.float32 = pa.Field() + + class Config: + batch_size = (32,) + +# Validate using the model - schema is built automatically +td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randint(0, 4, (32,)), "reward": torch.randn(32)}, + batch_size=[32], +) +RL.validate(td) +``` + +**Note:** Type annotations specify the dtype (e.g., `torch.float32`, `torch.int64`). +Use {func}`~pandera.tensordict.Field` to define additional constraints like shape and checks. + +## Guide contents + +```{toctree} +:maxdepth: 2 +:hidden: + +tensordict_schema +tensordict_model +tensordict_checks +error_reporting +``` + +- {ref}`pytorch-tensordict-schema` — validating a {class}`~tensordict.TensorDict` with `Tensor` components +- {ref}`pytorch-tensordict-model` — class-based `TensorDictModel` +- {ref}`pytorch-checks` — checks, parsers, and lazy validation +- {ref}`pytorch-error-reporting` — `SchemaError` / `SchemaErrors`, lazy validation, and failure cases + +## See also + +- {ref}`supported-dataframe-libraries` — other backends +- {ref}`checks` — general `Check` behaviour +- {ref}`lazy-validation` — `lazy=True` and `SchemaErrors` +- {ref}`configuration` — `ValidationDepth` and environment variables diff --git a/docs/source/pytorch_guide/tensordict_checks.md b/docs/source/pytorch_guide/tensordict_checks.md new file mode 100644 index 000000000..f06e6bd5b --- /dev/null +++ b/docs/source/pytorch_guide/tensordict_checks.md @@ -0,0 +1,73 @@ +--- +file_format: mystnb +--- + +(pytorch-checks)= + +# Checks and Lazy Validation + +Use {class}`~pandera.api.checks.Check` to validate tensor values. + +## Apply checks to Tensor components + +```{code-cell} python +import torch +from tensordict import TensorDict +from pandera import Check +import pandera.tensordict as pa + +schema = pa.TensorDictSchema( + keys={ + "values": pa.Tensor( + dtype=torch.float32, + shape=(None,), + checks=[Check.greater_than(0.0), Check.less_than(1.0)], + ), + }, + batch_size=(10,), +) + +td = TensorDict( + {"values": torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95])}, + batch_size=[10], +) +validated = schema.validate(td) +``` + +## Lazy validation + +Set `lazy=True` to collect all errors: + +```{code-cell} python +td_invalid = TensorDict( + {"values": torch.tensor([-0.1, 0.2, 0.3, 10.0, 0.5])}, + batch_size=[5], +) + +try: + schema.validate(td_invalid, lazy=True) +except pa.errors.SchemaErrors as exc: + print(f"Found {len(exc.schema_errors)} errors") + print(exc) +``` + +## Available checks + +Standard pandera checks work with tensor data: + +- {func}`~pandera.api.checks.Check.greater_than` +- {func}`~pandera.api.checks.Check.less_than` +- {func}`~pandera.api.checks.Check.greater_than_or_equal_to` +- {func}`~pandera.api.checks.Check.less_than_or_equal_to` +- {func}`~pandera.api.checks.Check.in_range` +- {func}`~pandera.api.checks.Check.isin` +- {func}`~pandera.api.checks.Check.notin` +- {func}`~pandera.api.checks.Check.str_matches` +- {func}`~pandera.api.checks.Check.str_contains` +- {func}`~pandera.api.checks.Check.str_startswith` +- {func}`~pandera.api.checks.Check.str_endswith` + +## See also + +- {ref}`checks` — full Check API reference +- {ref}`pytorch-error-reporting` — error handling diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md new file mode 100644 index 000000000..4c0c93234 --- /dev/null +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -0,0 +1,123 @@ +--- +file_format: mystnb +--- + +(pytorch-tensordict-model)= + +# TensorDictModel + +Use {class}`~pandera.api.tensordict.model.TensorDictModel` to define a class-based +schema that maps directly to a {class}`~tensordict.TensorDict`. + +## Define a model + +```{code-cell} python +import torch +import pandera.tensordict as pa + +class RL(pa.TensorDictModel): + """Schema for reinforcement learning data.""" + + # Type annotation specifies the dtype (torch.float32, torch.int64, etc.) + observation: torch.float32 = pa.Field(shape=(None, 10)) + action: torch.int64 = pa.Field(shape=(None,)) + reward: torch.float32 = pa.Field() +``` + +Use PyTorch dtypes in type annotations (e.g., `torch.float32`, `torch.int64`) +to specify the expected data type. Use {func}`~pandera.tensordict.Field` to +define additional constraints. + +## Validate with a model + +```{code-cell} python +from tensordict import TensorDict + +td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randint(0, 4, (32,)), "reward": torch.randn(32)}, + batch_size=[32], +) +validated = RL.validate(td) +``` + +## Field configuration + +Use {func}`~pandera.tensordict.Field` to customize field options: + +- `shape`: Expected shape tuple (use `None` for variable dimensions) +- `checks`: List of Check instances or check arguments +- `nullable`: Whether the key can be missing +- `default`: Default value if missing + +```{code-cell} python +class RLWithConfig(pa.TensorDictModel): + """RL schema with field-level checks.""" + + observation: torch.float32 = pa.Field( + shape=(None, 10), + ge=-1.0, + le=1.0, + ) + action: torch.int64 = pa.Field( + shape=(None,), + isin=[0, 1, 2, 3], + ) + reward: torch.float32 = pa.Field( + gt=0.0, + ) + + class Config: + batch_size = (32,) +``` + +## Configuring the model + +Use a nested `Config` class to configure schema-level options: + +- `batch_size`: Expected batch size tuple (use `None` for variable dimensions) + +```{code-cell} python +class RLWithBatchSize(pa.TensorDictModel): + observation: torch.float32 = pa.Field(shape=(None, 10)) + action: torch.int64 = pa.Field(shape=(None,)) + + class Config: + batch_size = (64,) +``` + +## Lazy validation + +Use `lazy=True` to collect all validation errors: + +```{code-cell} python +try: + RL.validate(td, lazy=True) +except pa.SchemaErrors as e: + print(f"Found {len(e.schema_errors)} validation errors:") + for err in e.schema_errors: + print(f" - {err.reason_code}") +``` + +## Model inheritance + +Models can be inherited to create more specific schemas: + +```{code-cell} python +class BaseRL(pa.TensorDictModel): + observation: torch.float32 = pa.Field(shape=(None, 10)) + + class Config: + batch_size = (32,) + +class ExtendedRL(BaseRL): + action: torch.int64 = pa.Field(shape=(None,)) + +# ExtendedRL has both 'observation' and 'action' +schema = ExtendedRL.to_schema() +``` + +## See also + +- {ref}`pytorch-tensordict-schema` — dictionary-based schema +- {ref}`pytorch-checks` — value checks +- {ref}`configuration` — validation configuration diff --git a/docs/source/pytorch_guide/tensordict_schema.md b/docs/source/pytorch_guide/tensordict_schema.md new file mode 100644 index 000000000..56eaf8b5d --- /dev/null +++ b/docs/source/pytorch_guide/tensordict_schema.md @@ -0,0 +1,94 @@ +--- +file_format: mystnb +--- + +(pytorch-tensordict-schema)= + +# TensorDictSchema + +Use {class}`~pandera.api.tensordict.container.TensorDictSchema` to validate +{class}`~tensordict.TensorDict` objects. + +## Define a schema + +Define a dictionary-based schema with {class}`~pandera.api.tensordict.components.Tensor` +components: + +```{code-cell} python +import torch +from tensordict import TensorDict +import pandera.tensordict as pa + +schema = pa.TensorDictSchema( + keys={ + "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)), + "action": pa.Tensor(dtype=torch.float32, shape=(None, 5)), + }, + batch_size=(32,), +) +``` + +- `keys`: A dictionary mapping key names to {class}`~pandera.api.tensordict.components.Tensor` objects +- `batch_size`: Required batch dimension sizes. Use `None` for flexible dimensions. + +## Validate a TensorDict + +```{code-cell} python +td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + batch_size=[32], +) +validated = schema.validate(td) +``` + +## Lazy validation + +Set `lazy=True` to collect all errors instead of fail-fast. Note that TensorDict +validates batch dimensions on creation, so we need to use matching dimensions: + +```{code-cell} python +td_wrong_dtype = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randint(0, 5, (32,))}, + batch_size=[32], +) + +try: + schema.validate(td_wrong_dtype, lazy=True) +except pa.errors.SchemaErrors as exc: + print(f"Total errors: {len(exc.schema_errors)}") + for err in exc.schema_errors: + print(f" - {str(err)}") +``` + +## Tensor component + +The {class}`~pandera.api.tensordict.components.Tensor` component validates individual tensor keys. + +### Parameters + +- `dtype`: Torch dtype (e.g., `torch.float32`, `torch.int64`) +- `shape`: Expected shape tuple. Use `None` for flexible dimensions. +- `checks`: Optional list of {class}`~pandera.api.checks.Check` instances + +### Example + +```{code-cell} python +from pandera import Check + +schema = pa.TensorDictSchema( + keys={ + "values": pa.Tensor( + dtype=torch.float32, + shape=(None,), + checks=[Check.greater_than(0.0), Check.less_than(1.0)], + ), + }, + batch_size=(10,), +) +``` + +## See also + +- {ref}`pytorch-tensordict-model` — class-based schema definition +- {ref}`pytorch-checks` — checks for value validation +- {ref}`pytorch-error-reporting` — error handling diff --git a/docs/source/reference/index.md b/docs/source/reference/index.md index 3a510ce77..f48e93b72 100644 --- a/docs/source/reference/index.md +++ b/docs/source/reference/index.md @@ -66,6 +66,8 @@ See :ref:`api-core` for full details. * - :ref:`Xarray ` - Schemas for labelled N-dimensional :mod:`xarray` arrays, datasets, and datatrees + * - :ref:`PyTorch ` + - Schemas for :mod:`tensordict` TensorDict and tensorclass objects ``` ```{toctree} @@ -82,4 +84,5 @@ strategies extensions errors xarray +pytorch ``` diff --git a/docs/source/reference/pytorch.rst b/docs/source/reference/pytorch.rst new file mode 100644 index 000000000..491ac59a3 --- /dev/null +++ b/docs/source/reference/pytorch.rst @@ -0,0 +1,78 @@ +.. _api-pytorch: + +PyTorch +======= + +*New in 0.32.0* + +Schemas and components for validating :class:`tensordict.TensorDict` and +:class:`tensordict.tensorclass`. Typical imports use :mod:`pandera.tensordict`; +implementations live under :mod:`pandera.api.tensordict`. +See :ref:`pytorch-guide` for usage examples. + +The :mod:`pandera.tensordict` entry point also re-exports +:class:`~pandera.api.checks.Check` and :mod:`pandera.errors` — +see :ref:`api-core` for those APIs. + +Schemas +------- + +.. autosummary:: + :toctree: generated + :template: class.rst + :nosignatures: + + pandera.api.tensordict.container.TensorDictSchema + +Schema components +----------------- + +.. autosummary:: + :toctree: generated + :template: class.rst + :nosignatures: + + pandera.api.tensordict.components.Tensor + +Declarative models +------------------ + +.. autosummary:: + :toctree: generated + :template: class.rst + :nosignatures: + + pandera.api.tensordict.model.TensorDictModel + +.. autosummary:: + :toctree: generated + :nosignatures: + + pandera.api.tensordict.model_components.Field + +Configuration +------------- + +See also :ref:`api-core` for the full configuration API. + +.. autosummary:: + :toctree: generated + :template: class.rst + :nosignatures: + + pandera.config.ValidationDepth + pandera.config.ValidationScope + +Dtypes +------ + +See :ref:`api-dtypes` for :class:`~pandera.engines.tensordict_engine.DataType` and +:class:`~pandera.engines.tensordict_engine.Engine`. + +See also +-------- + +- :ref:`pytorch-guide` — tutorials and examples +- :ref:`api-core` — :class:`~pandera.api.checks.Check`, configuration, errors +- :ref:`api-dtypes` — dtype engines including PyTorch +- :ref:`api-decorators` — validation decorators diff --git a/environment.yml b/environment.yml index e2d08a3a0..25cf70c45 100644 --- a/environment.yml +++ b/environment.yml @@ -33,6 +33,10 @@ dependencies: # xarray extra - xarray >= 2024.10.0 + # torch extra + - torch + - tensordict + # modin extra - modin - protobuf diff --git a/noxfile.py b/noxfile.py index 0df160fdc..9d0038758 100644 --- a/noxfile.py +++ b/noxfile.py @@ -37,6 +37,12 @@ ] ) +EXTRAS_REQUIRING_TORCH = frozenset( + [ + "torch", + ] +) + CI_RUN = os.environ.get("CI") == "true" if CI_RUN: print("Running on CI") @@ -139,6 +145,12 @@ def _testing_requirements( PYPROJECT["project"]["optional-dependencies"]["pandas"] ) + # torch extra requires torch and tensordict + if extra in EXTRAS_REQUIRING_TORCH: + _requirements.extend( + PYPROJECT["project"]["optional-dependencies"]["torch"] + ) + _requirements = list(set(_requirements)) _numpy: str | None = None @@ -199,6 +211,7 @@ def _testing_requirements( "dask", "ibis", "xarray", + "torch", } for extra in OPTIONAL_DEPENDENCIES: if extra == "pandas": @@ -298,12 +311,10 @@ def docs(session: Session) -> None: """Build the documentation.""" # this is needed until ray and geopandas are supported on python 3.10 - session.install("-e", ".") session.install( - *_testing_requirements( - session, extra="all", pandas=PANDAS_VERSIONS[0] - ), - *nox.project.dependency_groups(PYPROJECT, "dev", "testing", "docs"), + "-e", + ".[all]", + *nox.project.dependency_groups(PYPROJECT, "docs"), ) session.run("uv", "pip", "list") session.chdir("docs") @@ -341,6 +352,9 @@ def docs(session: Session) -> None: *args, ) + # Ensure torch is available for TensorDictModel doctests + session.run("python", "-c", "import torch") + session.run("xdoctest", PACKAGE, "--quiet") diff --git a/pandera/api/base/checks.py b/pandera/api/base/checks.py index c7cdf68f0..38b88ce33 100644 --- a/pandera/api/base/checks.py +++ b/pandera/api/base/checks.py @@ -3,14 +3,7 @@ import inspect from collections.abc import Callable, Iterable from itertools import chain -from typing import ( - Any, - NamedTuple, - Optional, - TypeVar, - Union, - no_type_check, -) +from typing import Any, NamedTuple, Optional, TypeVar, Union, no_type_check from pandera.api.function_dispatch import Dispatcher from pandera.backends.base import BaseCheckBackend diff --git a/pandera/api/base/model.py b/pandera/api/base/model.py index b957553be..035d473f8 100644 --- a/pandera/api/base/model.py +++ b/pandera/api/base/model.py @@ -2,13 +2,7 @@ import os from collections.abc import Mapping -from typing import ( - Any, - ClassVar, - Optional, - TypeVar, - Union, -) +from typing import Any, ClassVar, Optional, TypeVar, Union from pandera.api.base.model_components import BaseFieldInfo from pandera.api.base.model_config import BaseModelConfig diff --git a/pandera/api/base/model_components.py b/pandera/api/base/model_components.py index d660c25a8..776c03ea9 100644 --- a/pandera/api/base/model_components.py +++ b/pandera/api/base/model_components.py @@ -1,12 +1,7 @@ """Model component base classes.""" from collections.abc import Callable, Iterable -from typing import ( - Any, - Optional, - Union, - cast, -) +from typing import Any, Optional, Union, cast from pandera.api.checks import Check from pandera.api.parsers import Parser diff --git a/pandera/api/checks.py b/pandera/api/checks.py index b569073c1..dca583aa6 100644 --- a/pandera/api/checks.py +++ b/pandera/api/checks.py @@ -2,12 +2,7 @@ import re from collections.abc import Callable, Iterable -from typing import ( - Any, - Optional, - TypeVar, - Union, -) +from typing import Any, Optional, TypeVar, Union from pandera import errors from pandera.api.base.checks import BaseCheck, CheckResult diff --git a/pandera/api/dataframe/container.py b/pandera/api/dataframe/container.py index 015598bd4..a2d5ae8c7 100644 --- a/pandera/api/dataframe/container.py +++ b/pandera/api/dataframe/container.py @@ -6,14 +6,7 @@ import sys import warnings from pathlib import Path -from typing import ( - Any, - Generic, - Optional, - TypeVar, - Union, - cast, -) +from typing import Any, Generic, Optional, TypeVar, Union, cast from pandera import errors from pandera.api.base.schema import BaseSchema diff --git a/pandera/api/dataframe/model_components.py b/pandera/api/dataframe/model_components.py index 213b71e13..cbe21062e 100644 --- a/pandera/api/dataframe/model_components.py +++ b/pandera/api/dataframe/model_components.py @@ -1,11 +1,7 @@ """DataFrameModel components""" from collections.abc import Callable, Iterable -from typing import ( - Any, - Union, - cast, -) +from typing import Any, Union, cast from pandera.api.base.model_components import ( BaseCheckInfo, diff --git a/pandera/api/ibis/model.py b/pandera/api/ibis/model.py index 05fe23948..d888a1018 100644 --- a/pandera/api/ibis/model.py +++ b/pandera/api/ibis/model.py @@ -2,10 +2,7 @@ import inspect import sys -from typing import ( - Optional, - cast, -) +from typing import Optional, cast import ibis import ibis.expr.datatypes as dt diff --git a/pandera/api/polars/model.py b/pandera/api/polars/model.py index 4bd7dd1f5..e0d439deb 100644 --- a/pandera/api/polars/model.py +++ b/pandera/api/polars/model.py @@ -245,5 +245,7 @@ def empty(cls: type[Self], *_args) -> DataFrame[Self]: """Create an empty DataFrame with the schema of this model.""" schema = copy.deepcopy(cls.to_schema()) schema.coerce = True - empty_df = schema.coerce_dtype(pl.DataFrame(schema=[*schema.columns])) - return DataFrame[Self](empty_df) + coerced = schema.coerce_dtype(pl.DataFrame(schema=[*schema.columns])) + if isinstance(coerced, pl.LazyFrame): + coerced = coerced.collect() + return DataFrame[Self](coerced) diff --git a/pandera/api/pyspark/components.py b/pandera/api/pyspark/components.py index f111d234a..9a5483f8e 100644 --- a/pandera/api/pyspark/components.py +++ b/pandera/api/pyspark/components.py @@ -9,11 +9,7 @@ from pandera.dtypes import DataType from pandera.engines import pyspark_engine -from .types import ( - PySparkDataFrameTypes, - PySparkDtypeInputTypes, - PySparkFrame, -) +from .types import PySparkDataFrameTypes, PySparkDtypeInputTypes, PySparkFrame class Column(ComponentSchema[PySparkDataFrameTypes]): diff --git a/pandera/api/pyspark/container.py b/pandera/api/pyspark/container.py index 14fce3178..8e49c50ec 100644 --- a/pandera/api/pyspark/container.py +++ b/pandera/api/pyspark/container.py @@ -16,11 +16,7 @@ from pandera.engines import pyspark_engine from pandera.utils import docstring_substitution -from .types import ( - PySparkDataFrameTypes, - PySparkDtypeInputTypes, - PySparkFrame, -) +from .types import PySparkDataFrameTypes, PySparkDtypeInputTypes, PySparkFrame if TYPE_CHECKING: import pandera.api.pyspark.components diff --git a/pandera/api/pyspark/model.py b/pandera/api/pyspark/model.py index 1acdef110..e39b0e29d 100644 --- a/pandera/api/pyspark/model.py +++ b/pandera/api/pyspark/model.py @@ -5,13 +5,7 @@ import re import typing from collections.abc import Callable, Iterable, Mapping -from typing import ( - Any, - Optional, - TypeVar, - Union, - cast, -) +from typing import Any, Optional, TypeVar, Union, cast from pyspark.sql.types import StructType from typing_extensions import Self diff --git a/pandera/api/pyspark/model_components.py b/pandera/api/pyspark/model_components.py index c7dce586c..fa5c73aca 100644 --- a/pandera/api/pyspark/model_components.py +++ b/pandera/api/pyspark/model_components.py @@ -1,11 +1,7 @@ """DataFrameModel components""" from collections.abc import Callable, Iterable -from typing import ( - Any, - TypeVar, - Union, -) +from typing import Any, TypeVar, Union from pandera.api.dataframe.components import ComponentSchema from pandera.api.dataframe.model_components import Field as _Field diff --git a/pandera/api/pyspark/types.py b/pandera/api/pyspark/types.py index f370b8f0b..2a44167ca 100644 --- a/pandera/api/pyspark/types.py +++ b/pandera/api/pyspark/types.py @@ -20,9 +20,7 @@ ) from pyspark.sql.connect.group import GroupedData else: - from pyspark.sql import ( - DataFrame as PySparkConnectDataFrame, - ) + from pyspark.sql import DataFrame as PySparkConnectDataFrame from pyspark.sql.group import GroupedData PySparkDataFrameTypes = Union[ diff --git a/pandera/api/tensordict/__init__.py b/pandera/api/tensordict/__init__.py new file mode 100644 index 000000000..f9332737b --- /dev/null +++ b/pandera/api/tensordict/__init__.py @@ -0,0 +1,27 @@ +"""Pandera TensorDict API.""" + +from __future__ import annotations + +# mypy: disable-error-code=attr-defined +from typing import TYPE_CHECKING + +from pandera import errors + +if TYPE_CHECKING: + from pandera.api.tensordict.components import Tensor + from pandera.api.tensordict.container import TensorDictSchema + from pandera.api.tensordict.model import TensorDictModel + +# Import actual implementations (not just for type checking) +from pandera.api.tensordict.components import Tensor +from pandera.api.tensordict.container import TensorDictSchema +from pandera.api.tensordict.model import TensorDictModel +from pandera.api.tensordict.model_components import Field + +__all__ = [ + "Tensor", + "TensorDictSchema", + "TensorDictModel", + "Field", + "errors", +] diff --git a/pandera/api/tensordict/components.py b/pandera/api/tensordict/components.py new file mode 100644 index 000000000..d793ce85a --- /dev/null +++ b/pandera/api/tensordict/components.py @@ -0,0 +1,65 @@ +"""Tensor component for TensorDict schemas.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional, Union + +from pandera.api.base.types import CheckList +from pandera.engines import tensordict_engine + + +class Tensor: + """Validate a single tensor entry in a TensorDict. + + Analogous to :class:`~pandera.api.pandas.components.Column` for DataFrames + and :class:`~pandera.api.xarray.components.DataVar` for xarray Datasets. + """ + + def __init__( + self, + dtype: Any = None, + shape: tuple[int | None, ...] | None = None, + checks: CheckList | None = None, + nullable: bool = False, + coerce: bool = False, + name: str | None = None, + required: bool = True, + ) -> None: + """Create tensor validator. + + :param dtype: PyTorch dtype (e.g., ``torch.float32``, ``torch.int64``). + If a string is provided, it will be converted to the corresponding + torch.dtype. + :param shape: Expected shape tuple. Use ``None`` for flexible dimensions. + For example, ``(None, 10)`` allows any batch size with 10 features. + :param checks: List of :class:`~pandera.api.checks.Check` instances for + value validation. + :param nullable: Whether the key can be missing from the TensorDict. + :param coerce: Whether to coerce dtype during validation. + :param name: Tensor key name in TensorDict. If ``None``, uses the key + from the schema definition. + :param required: Whether the tensor is required (always ``True`` for + now, as missing keys are handled separately). + """ + self.dtype = ( + tensordict_engine.Engine.dtype(dtype) + if dtype is not None + else None + ) + self.shape = shape + + # Normalize checks to a list + self.checks: list = [] + if checks is not None: + if isinstance(checks, (list, tuple)): + self.checks = list(checks) + else: + self.checks = [checks] + + self.nullable = nullable + self.coerce = coerce + self.name = name + self.required = required + + def __repr__(self) -> str: + return f"Tensor(dtype={self.dtype}, shape={self.shape})" diff --git a/pandera/api/tensordict/container.py b/pandera/api/tensordict/container.py new file mode 100644 index 000000000..ea260cc9c --- /dev/null +++ b/pandera/api/tensordict/container.py @@ -0,0 +1,157 @@ +"""TensorDict container specification.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional + +from pandera.api.base.schema import BaseSchema +from pandera.api.tensordict.components import Tensor +from pandera.backends.base import BaseSchemaBackend + + +class TensorDictSchema(BaseSchema): + """Validate TensorDict and tensorclass objects. + + Analogous to :class:`~pandera.api.pandas.DataFrameSchema` for DataFrames + and :class:`~pandera.api.xarray.DatasetSchema` for xarray Datasets. + + The schema accepts a ``keys`` parameter (instead of ``columns``) which is + more semantically appropriate for TensorDict's dictionary-like structure. + """ + + def __init__( + self, + keys: dict[str, Tensor] | None = None, + batch_size: tuple[int | None, ...] | None = None, + dtype: Any = None, + checks=None, + coerce: bool = False, + name: str | None = None, + title: str | None = None, + description: str | None = None, + metadata: dict | None = None, + ) -> None: + """Create TensorDict schema. + + :param keys: Dictionary mapping key names to :class:`Tensor` components. + If a list of strings is provided, creates default Tensor components + for each key. + :param batch_size: Expected batch dimensions. Use ``None`` for flexible + dimensions. For example, ``(32,)`` requires exactly 32 samples, + while ``(None,)`` allows any batch size. + :param dtype: Override dtype for all tensors (optional). + :param checks: Schema-wide checks applied to the entire TensorDict. + :param coerce: Enable automatic dtype coercion during validation. + :param name: Name of the schema. + :param title: Human-readable title for the schema. + :param description: Detailed description of the schema. + :param metadata: Additional metadata associated with the schema. + """ + if keys is None: + keys = {} + elif isinstance(keys, list): + keys = {k: Tensor() for k in keys} + + super().__init__( + dtype=dtype, + checks=checks, + coerce=coerce, + name=name, + title=title, + description=description, + metadata=metadata, + ) + + self.keys: dict[str, Tensor] = keys + self.batch_size = batch_size + + @staticmethod + def register_default_backends(check_obj_cls: type) -> None: + """Register default backends at validation time.""" + from pandera.backends.tensordict.register import ( + register_tensordict_backends, + ) + + register_tensordict_backends() + + @classmethod + def get_backend( + cls, check_obj: Any | None = None, check_type: type | None = None + ) -> BaseSchemaBackend: + """Get the backend for TensorDict validation. + + Overrides base class to handle tensorclass objects which have unique + types at runtime. We use TensorDict's backend since tensorclasses share + the same interface and both inherit from object in their MRO. + """ + if check_obj is not None: + try: + # Check if it's a valid TensorDict/tensorclass object + from pandera.backends.tensordict.base import _is_tensordict + + if _is_tensordict(check_obj): + check_type = type(check_obj) + cls.register_default_backends(check_type) + + # Try to get backend for this exact type first + classes = __import__("inspect").getmro(check_type) + for _class in classes: + try: + return cls.BACKEND_REGISTRY[(cls, _class)]() + except KeyError: + pass + + # For tensorclass instances (which don't have a specific backend), + # fall back to TensorDict's backend + from tensordict import TensorDict as TD + + if check_type != TD and hasattr( + check_obj, "_is_tensorclass" + ): + try: + return cls.BACKEND_REGISTRY[(cls, TD)]() + except KeyError: + pass + except ImportError: + pass + + # Default fallback to base class implementation + return super().get_backend(check_obj) + + def validate( + self, + check_obj: Any, + head: int | None = None, + tail: int | None = None, + sample: int | None = None, + random_state: int | None = None, + lazy: bool = False, + inplace: bool = False, + ) -> Any: + """Validate a TensorDict against the schema. + + :param check_obj: TensorDict or tensorclass to validate. + :param head: Number of leading samples to validate (not used for TensorDict). + :param tail: Number of trailing samples to validate (not used for TensorDict). + :param sample: Number of random samples to validate (not used for TensorDict). + :param random_state: Random seed for sampling (not used for TensorDict). + :param lazy: If ``True``, collect all errors instead of fail-fast. + :param inplace: Whether to modify object in place (not supported for TensorDict). + :returns: Validated TensorDict or tensorclass. + :raises SchemaError: If validation fails. + """ + return self.get_backend(check_obj).validate( + check_obj=check_obj, + schema=self, + head=head, + tail=tail, + sample=sample, + random_state=random_state, + lazy=lazy, + inplace=inplace, + ) + + def __repr__(self) -> str: + """Return string representation of the schema.""" + keys_str = ", ".join(repr(k) for k in self.keys.keys()) + return f"TensorDictSchema(keys={{{keys_str}}}, batch_size={self.batch_size})" diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py new file mode 100644 index 000000000..dadd6ced9 --- /dev/null +++ b/pandera/api/tensordict/model.py @@ -0,0 +1,646 @@ +"""TensorDictModel for class-based schema definitions.""" + +from __future__ import annotations + +import inspect +import sys +import threading +import types as types_module +import typing +from typing import Any, ClassVar, cast + +import typing_inspect + +# Import torch if available for TensorDict type hints in docstrings +try: + import torch # noqa: F401 +except ImportError: + pass + +from pandera.api.base.model import BaseModel +from pandera.api.checks import Check +from pandera.api.parsers import Parser +from pandera.api.tensordict.components import Tensor +from pandera.api.tensordict.container import TensorDictSchema +from pandera.api.tensordict.model_config import BaseConfig +from pandera.errors import SchemaInitError +from pandera.typing import AnnotationInfo + +TFields = dict[str, tuple[AnnotationInfo, Any]] +TChecks = dict[str, list[Check]] +TParsers = dict[str, list[Parser]] + +_CONFIG_KEY = "Config" +MODEL_CACHE: dict[tuple[type[TensorDictModel], int], Any] = {} + + +def _is_field(name: str) -> bool: + return not name.startswith("_") and name != _CONFIG_KEY + + +def _metaclass_dict(cls: type) -> Any: + """``type(cls).__dict__`` (typeshed types ``.__dict__`` on metaclass poorly).""" + return cast(Any, type(cls).__dict__) + + +def _get_tensor_dict_class_type_hints( + cls: type, + globalns: dict[str, Any] | None = None, + localns: dict[str, Any] | None = None, + *, + include_extras: bool = True, +) -> dict[str, Any]: + """Resolve class annotations like :func:`typing.get_type_hints`. + + String annotations (PEP 563) such as ``"torch.float32"`` evaluate to + ``torch.dtype`` instances; :func:`typing.get_type_hints` rejects those + because they are not :class:`type` objects. TensorDict field annotations + intentionally use dtypes, so string annotations are resolved with + :func:`eval` like :func:`inspect.get_annotations` (same security model as + normal forward references). + """ + # typing.get_type_hints uses these; they are not in typeshed (private API). + _eval_type = cast(Any, getattr(typing, "_eval_type")) + _strip_annotations = cast(Any, getattr(typing, "_strip_annotations")) + + if getattr(cls, "__no_type_check__", None): + return {} + hints: dict[str, Any] = {} + for base in reversed(cls.__mro__): + if globalns is None: + base_globals = getattr( + sys.modules.get(base.__module__, None), + "__dict__", + {}, + ) + else: + base_globals = globalns + ann = base.__dict__.get("__annotations__", {}) + if isinstance(ann, types_module.GetSetDescriptorType): + ann = {} + base_locals = dict(vars(base)) if localns is None else localns + if localns is None and globalns is None: + base_globals, base_locals = base_locals, base_globals + for name, value in ann.items(): + if value is None: + value = type(None) + if isinstance(value, str): + value = eval(value, base_globals, base_locals) + else: + value = _eval_type(value, base_globals, base_locals) + hints[name] = value + if include_extras: + return hints + return {k: _strip_annotations(t) for k, t in hints.items()} + + +class _FieldsDescriptor: + def __init__(self) -> None: + self.cache: dict[type, Any] = {} + + def __get__(self, obj: Any, cls: type[TensorDictModel]) -> TFields: + if self.cache.get(cls) is None: + self.cache[cls] = cls._collect_fields() + return self.cache[cls] + + +class _SchemaDescriptor: + def __get__(self, obj: Any, cls: type[TensorDictModel]) -> Any: + tid = threading.get_ident() + key = (cls, tid) + if MODEL_CACHE.get(key) is None: + try: + MODEL_CACHE[key] = cls.build_schema_() + except Exception as e: + raise AttributeError( + f"'{cls.__name__}' does not implement build_schema_() and cannot " + f"generate a schema. To be able to generate a schema, subclass the" + "TensorDictModel." + ) from e + return MODEL_CACHE[key] + + +class _ChecksDescriptor: + def __init__(self) -> None: + self.cache: dict[type, Any] = {} + + def __get__(self, obj: Any, cls: type[TensorDictModel]) -> TChecks: + if self.cache.get(cls) is None: + check_infos = typing.cast( + list[Any], cls._collect_check_infos(CHECK_KEY) + ) + self.cache[cls] = cls._extract_checks( + check_infos, field_names=list(cls.__fields__.keys()) + ) + return self.cache[cls] + + +class _RootCheckDescriptor: + def __init__(self) -> None: + self.cache: dict[type, Any] = {} + + def __get__(self, obj: Any, cls: type[TensorDictModel]) -> list[Check]: + if self.cache.get(cls) is None: + df_check_infos = cls._collect_check_infos(DATAFRAME_CHECK_KEY) + custom = cls._extract_df_checks(df_check_infos) + reg = _convert_extras_to_checks( + {} if cls.__extras__ is None else cls.__extras__ + ) + self.cache[cls] = custom + reg + return self.cache[cls] + + +class _ParsersDescriptor: + def __init__(self) -> None: + self.cache: dict[type, Any] = {} + + def __get__(self, obj: Any, cls: type[TensorDictModel]) -> TParsers: + if self.cache.get(cls) is None: + parser_infos = typing.cast( + list[Any], cls._collect_parser_infos(PARSER_KEY) + ) + self.cache[cls] = cls._extract_parsers( + parser_infos, field_names=list(cls.__fields__.keys()) + ) + return self.cache[cls] + + +class _RootParsersDescriptor: + def __init__(self) -> None: + self.cache: dict[type, Any] = {} + + def __get__(self, obj: Any, cls: type[TensorDictModel]) -> list[Parser]: + if self.cache.get(cls) is None: + df_parser_infos = cls._collect_parser_infos(DATAFRAME_PARSER_KEY) + self.cache[cls] = cls._extract_df_parsers(df_parser_infos) + return self.cache[cls] + + +CHECK_KEY = "__check_config__" +DATAFRAME_CHECK_KEY = "__dataframe_check_config__" +PARSER_KEY = "__parser_config__" +DATAFRAME_PARSER_KEY = "__dataframe_parser_config__" + + +def _convert_extras_to_checks(extras: dict[str, Any]) -> list[Check]: + """Convert extras to checks.""" + from pandera.api.checks import Check + + checks = [] + for name, value in extras.items(): + if isinstance(value, tuple): + args, kwargs = value, {} + elif isinstance(value, dict): + args, kwargs = (), value + elif value is Ellipsis: + args, kwargs = (), {} + else: + args, kwargs = (value,), {} + + checks.append(getattr(Check, name)(*args, **kwargs)) + + return checks + + +class TensorDictModel(BaseModel): + """Declarative TensorDict schema using class definitions. + + Fields are defined using ``pa.Field`` directly in type annotations. Use PyTorch dtypes + (e.g., `torch.float32`, `torch.int64`) to specify the expected data type: + + Example: + >>> import torch + >>> import pandera.tensordict as pa + + >>> class MySchema(pa.TensorDictModel): + ... observation: torch.float32 = pa.Field(shape=(None, 10)) + ... action: torch.int64 = pa.Field(shape=(None,)) + ... + ... class Config: + ... batch_size = (32,) + + >>> schema = MySchema.to_schema() + """ + + Config: type[BaseConfig] = BaseConfig + __schema__ = _SchemaDescriptor() + __config__: type[BaseConfig] | None = None + + #: Key according to `FieldInfo.name` + __fields__: ClassVar[TFields] = cast(TFields, _FieldsDescriptor()) + __checks__ = cast(TChecks, _ChecksDescriptor()) + __parsers__: TParsers = cast(TParsers, _ParsersDescriptor()) + __root_checks__ = cast(list[Check], _RootCheckDescriptor()) + __root_parsers__ = cast(list[Parser], _RootParsersDescriptor()) + + def __init_subclass__(cls, **kwargs: Any) -> None: + """Initialize subclass and build schema from class annotations.""" + if _CONFIG_KEY in cls.__dict__: + if not hasattr(cls.Config, "name"): + cls.Config.name = None + else: + cls.Config = type("Config", (cls.Config,), {}) + + super().__init_subclass__(**kwargs) + hints = _get_tensor_dict_class_type_hints(cls, include_extras=True) + + for fname in hints.keys(): + if not _is_field(fname): + continue + if fname in cls.__dict__: + continue + # Auto-add Field() if not explicitly defined + from pandera.api.tensordict.model_components import Field + + field = Field() + field.__set_name__(cls, fname) + setattr(cls, fname, field) + + cls.__config__, cls.__extras__ = cls._collect_config_and_extras() + + @classmethod + def _get_model_attrs(cls) -> dict[str, Any]: + """Return all attributes from the model class.""" + bases = inspect.getmro(cls)[:-1] + attrs: dict[str, Any] = {} + for base in reversed(bases): + if ( + issubclass(base, TensorDictModel) + and base is not TensorDictModel + ): + attrs.update(base.__dict__) + return attrs + + @classmethod + def _collect_fields(cls) -> TFields: + """Centralize publicly named fields and their corresponding annotations.""" + from pandera.api.tensordict.model_components import TensorDictFieldInfo + + annotations = _get_tensor_dict_class_type_hints( + cls, include_extras=True + ) + attrs = cls._get_model_attrs() + + missing = [] + for name, attr in attrs.items(): + if inspect.isroutine(attr): + continue + if not _is_field(name): + annotations.pop(name, None) + elif name not in annotations: + missing.append(name) + + if missing: + raise SchemaInitError(f"Found missing annotations: {missing}") + + fields = {} + for field_name, annotation in annotations.items(): + if not _is_field(field_name): + continue + field = attrs[field_name] + + # Check if it's a TensorDictFieldInfo instance + if not isinstance(field, TensorDictFieldInfo): + raise SchemaInitError( + f"'{field_name}' can only be assigned a 'Field', " + f"not a '{type(field)}'." + ) + fields[field.name] = (AnnotationInfo(annotation), field) + return fields + + @classmethod + def _extract_config_options_and_extras( + cls, + config: Any, + ) -> tuple[dict[str, Any], dict[str, Any]]: + """Extract config options and extras from Config class.""" + from pandera.api.dataframe.model_config import ( + BaseConfig as DataFrameBaseConfig, + ) + + config_options = {} + extras = {} + + # Get allowed config option names + if hasattr(DataFrameBaseConfig, "__dataclass_fields__"): + config_field_names = list( + DataFrameBaseConfig.__dataclass_fields__.keys() + ) + else: + config_field_names = [ + attr + for attr in vars(cls.Config) + if not attr.startswith("_") + and callable(getattr(cls.Config, attr)) is False + ] + + for name, value in vars(config).items(): + if name in config_field_names: + config_options[name] = value + elif _is_field(name): + extras[name] = value + + return config_options, extras + + @classmethod + def _collect_config_and_extras( + cls, + ) -> tuple[type[BaseConfig], dict[str, Any]]: + """Collect config options from bases, splitting off unknown options.""" + bases = inspect.getmro(cls)[:-1] + bases = tuple( + base for base in bases if issubclass(base, TensorDictModel) + ) + root_model, *models = reversed(bases) + + options, extras = cls._extract_config_options_and_extras( + root_model.Config + ) + + for model in models: + config = getattr(model, _CONFIG_KEY, {}) + base_options, base_extras = cls._extract_config_options_and_extras( + config + ) + options.update(base_options) + extras.update(base_extras) + + return type("Config", (cls.Config,), options), extras + + @classmethod + def _collect_check_infos(cls, key: str) -> list[Any]: + """Collect inherited check metadata from bases.""" + bases = inspect.getmro(cls)[:-2] + bases = tuple( + base for base in bases if issubclass(base, TensorDictModel) + ) + + method_names = set() + check_infos = [] + for base in bases: + for attr_name, attr_value in vars(base).items(): + check_info = getattr(attr_value, key, None) + if not isinstance( + check_info, + _metaclass_dict(cls).get("Field", object).__bases__[0] + if hasattr( + _metaclass_dict(cls).get("Field", object), + "__bases__", + ) + else object, + ): + continue + if attr_name in method_names: + continue + method_names.add(attr_name) + check_infos.append(check_info) + return check_infos + + @classmethod + def _collect_parser_infos(cls, key: str) -> list[Any]: + """Collect inherited parser metadata from bases.""" + bases = inspect.getmro(cls)[:-2] + bases = tuple( + base for base in bases if issubclass(base, TensorDictModel) + ) + + method_names = set() + parser_infos = [] + for base in bases: + for attr_name, attr_value in vars(base).items(): + parser_info = getattr(attr_value, key, None) + if not isinstance( + parser_info, + _metaclass_dict(cls).get("Field", object).__bases__[0] + if hasattr( + _metaclass_dict(cls).get("Field", object), + "__bases__", + ) + else object, + ): + continue + method_names.add(attr_name) + parser_infos.append(parser_info) + return parser_infos + + @staticmethod + def _regex_filter(seq: Any, regexps: Any) -> set[str]: + """Filter items matching at least one of the regexes.""" + import re + + matched: set[str] = set() + for pattern_str in regexps: + pattern = re.compile(pattern_str) + matched.update(filter(pattern.match, seq)) + return matched + + @classmethod + def _extract_checks( + cls, check_infos: list[Any], field_names: list[str] + ) -> dict[str, list[Check]]: + """Collect field annotations from bases in mro reverse order.""" + checks: dict[str, list[Check]] = {} + for check_info in check_infos: + if not hasattr(check_info, "fields"): + continue + + check_info_fields = { + field.name + if isinstance( + field, + _metaclass_dict(cls).get("Field", object).__bases__[0] + if hasattr( + _metaclass_dict(cls).get("Field", object), + "__bases__", + ) + else object, + ) + else field + for field in check_info.fields + } + + matched = ( + cls._regex_filter(field_names, check_info_fields) + if hasattr(check_info, "regex") and check_info.regex + else check_info_fields + ) + + check_ = ( + check_info.to_check(cls) + if hasattr(check_info, "to_check") + else None + ) + + for field in matched: + if field not in field_names: + raise SchemaInitError( + f"Check {check_.name if check_ else 'unknown'} is assigned to a non-existing field '{field}'." + ) + if field not in checks: + checks[field] = [] + if check_ is not None: + checks[field].append(check_) + return checks + + @classmethod + def _extract_df_checks(cls, check_infos: list[Any]) -> list[Check]: + """Collect dataframe-level checks.""" + result: list[Check] = [] + for check_info in check_infos: + if hasattr(check_info, "to_check"): + check_ = check_info.to_check(cls) + if check_ is not None: + result.append(check_) + return result + + @classmethod + def _extract_parsers( + cls, parser_infos: list[Any], field_names: list[str] + ) -> dict[str, list[Parser]]: + """Collect field parsers from bases in mro reverse order.""" + parsers: dict[str, list[Parser]] = {} + for parser_info in parser_infos: + if not hasattr(parser_info, "fields"): + continue + + parser_info_fields = { + field.name + if isinstance( + field, + _metaclass_dict(cls).get("Field", object).__bases__[0] + if hasattr( + _metaclass_dict(cls).get("Field", object), + "__bases__", + ) + else object, + ) + else field + for field in parser_info.fields + } + + matched = ( + cls._regex_filter(field_names, parser_info_fields) + if hasattr(parser_info, "regex") and parser_info.regex + else parser_info_fields + ) + + parser_ = ( + parser_info.to_parser(cls) + if hasattr(parser_info, "to_parser") + else None + ) + + for field in matched: + if field not in field_names: + raise SchemaInitError( + f"Parser {parser_.name if parser_ else 'unknown'} is assigned to a non-existing field '{field}'." + ) + if field not in parsers: + parsers[field] = [] + if parser_ is not None: + parsers[field].append(parser_) + return parsers + + @classmethod + def _extract_df_parsers(cls, parser_infos: list[Any]) -> list[Parser]: + """Collect dataframe-level parsers.""" + result: list[Parser] = [] + for parser_info in parser_infos: + if hasattr(parser_info, "to_parser"): + parser_ = parser_info.to_parser(cls) + if parser_ is not None: + result.append(parser_) + return result + + @classmethod + def build_schema_(cls, **kwargs) -> TensorDictSchema: + """Create TensorDictSchema from model class. + + :returns: TensorDictSchema with fields converted to Tensor components. + """ + cfg = cls.__config__ + fields = cls.__fields__ + + columns: dict[str, Any] = {} + + for fname, (ann, fi) in fields.items(): + # Get dtype from annotation + dtype = ann.arg + + # Handle DataType special case - convert to torch.dtype if possible + if dtype is not None: + try: + from pandera.engines.tensordict_engine import ( + DataType as _DataType, + ) + + if isinstance(dtype, type) and issubclass( + dtype, _DataType + ): + dtype = dtype.type + except Exception: + pass + + # Get shape from Field info + shape = fi.shape if hasattr(fi, "shape") else None + + # Build tensor kwargs + tensor_kwargs: dict[str, Any] = { + "dtype": dtype, + "shape": shape, + "name": fname, + } + + if hasattr(fi, "checks") and fi.checks: + tensor_kwargs["checks"] = fi.checks + + columns[fname] = Tensor(**tensor_kwargs) + + # Get batch_size from Config + batch_size: tuple[int | None, ...] | None = ( + getattr(cfg, "batch_size", None) if cfg else None + ) + + try: + return TensorDictSchema( + keys=columns or None, batch_size=batch_size + ) + except ImportError as e: + raise RuntimeError("Could not import TensorDictSchema") from e + + @classmethod + def to_schema(cls: type[TensorDictModel]) -> TensorDictSchema: + """Create :class:`~pandera.TensorDictSchema` from the :class:`.TensorDictModel`.""" + return cls.__schema__ + + @classmethod + def validate( + cls: type[TensorDictModel], + check_obj: Any, + head: int | None = None, + tail: int | None = None, + sample: int | None = None, + random_state: int | None = None, + lazy: bool = False, + inplace: bool = False, + ) -> Any: + """Validate data against model schema. + + :param check_obj: TensorDict or tensorclass to validate. + :returns: Validated data object. + """ + if cls.__schema__ is not None: + return cls.__schema__.validate( + check_obj=check_obj, + head=head, + tail=tail, + sample=sample, + random_state=random_state, + lazy=lazy, + inplace=inplace, + ) + raise RuntimeError("Model schema was not initialized") + + def __new__(cls, *args: Any, **kwargs: Any) -> Any: + """Validate data against model schema.""" + return cls.validate(*args, **kwargs) diff --git a/pandera/api/tensordict/model_components.py b/pandera/api/tensordict/model_components.py new file mode 100644 index 000000000..b83850740 --- /dev/null +++ b/pandera/api/tensordict/model_components.py @@ -0,0 +1,159 @@ +"""Field descriptors for TensorDict declarative models.""" + +from collections.abc import Iterable +from typing import Any, Union + +from pandera.api.base.model_components import to_checklist +from pandera.api.checks import Check +from pandera.api.dataframe.model_components import FieldInfo, _check_dispatch +from pandera.errors import SchemaInitError + +CHECK_KEY = "__check_config__" +TENSOR_CHECK_KEY = "__tensor_check_config__" + + +class TensorDictFieldInfo(FieldInfo): + """Field metadata for :class:`TensorDictModel`.""" + + def __init__( + self, + *, + shape: tuple[int | None, ...] | None = None, + required: bool = True, + checks: Any = None, + parses: Any = None, + nullable: bool = False, + unique: bool = False, + coerce: bool = False, + regex: bool = False, + alias: Any = None, + check_name: bool | None = None, + dtype_kwargs: dict[str, Any] | None = None, + title: str | None = None, + description: str | None = None, + default: Any | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + super().__init__( + checks=checks, + parses=parses, + nullable=nullable, + unique=unique, + coerce=coerce, + regex=regex, + alias=alias, + check_name=check_name, + dtype_kwargs=dtype_kwargs, + title=title, + description=description, + default=default, + metadata=metadata, + ) + self.shape = shape + self.required = required + + def to_tensor_kwargs( + self, + dtype: Any, + *, + optional: bool, + checks: Any = None, + ) -> dict[str, Any]: + """Keyword args for :class:`~pandera.api.tensordict.components.Tensor`.""" + return { + "dtype": dtype, + "checks": to_checklist(self.checks) + to_checklist(checks), + "shape": self.shape, + "name": self.alias or self.name, + "title": self.title, + "description": self.description, + "nullable": self.nullable, + "coerce": self.coerce, + } + + +def Field( + *, + dtype: Any = None, + shape: tuple[int | None, ...] | None = None, + eq: Any | None = None, + ne: Any | None = None, + gt: Any | None = None, + ge: Any | None = None, + lt: Any | None = None, + le: Any | None = None, + in_range: Union[ + tuple[Any, Any], + tuple[Any, Any, bool, bool], + tuple[Any, Any, bool, bool, bool], + tuple[Any, Any, bool, bool, bool, bool], + dict[str, Any], + None, + ] = None, + isin: Iterable[Any] | None = None, + notin: Iterable[Any] | None = None, + no_nan: bool = False, + no_inf: bool = False, + nullable: bool = False, + unique: bool = False, + coerce: bool = False, + regex: bool = False, + ignore_na: bool = True, + raise_warning: bool = False, + n_failure_cases: int | None = None, + alias: Any | None = None, + check_name: bool | None = None, + dtype_kwargs: dict[str, Any] | None = None, + title: str | None = None, + description: str | None = None, + default: Any | None = None, + metadata: dict[str, Any] | None = None, + required: bool = True, + **kwargs: Any, +) -> Any: + """Field specification for TensorDict models.""" + check_kwargs = { + "ignore_na": ignore_na, + "raise_warning": raise_warning, + "n_failure_cases": n_failure_cases, + } + args = locals() + checks = [] + + check_dispatch = _check_dispatch() + for key in kwargs: + if key not in check_dispatch: + raise SchemaInitError( + f"custom check '{key}' is not available. Make sure you use " + "pandera.extensions.register_check_method decorator to " + "register your custom check method." + ) + + for arg_name, check_constructor in check_dispatch.items(): + arg_value = args.get(arg_name, kwargs.get(arg_name)) + if arg_value is None: + continue + if isinstance(arg_value, dict): + check_ = check_constructor(**arg_value, **check_kwargs) + elif isinstance(arg_value, tuple): + check_ = check_constructor(*arg_value, **check_kwargs) + else: + check_ = check_constructor(arg_value, **check_kwargs) + checks.append(check_) + + return TensorDictFieldInfo( + checks=checks or None, + nullable=nullable, + unique=unique, + coerce=coerce, + regex=regex, + check_name=check_name, + alias=alias, + title=title, + description=description, + default=default, + dtype_kwargs=dtype_kwargs, + metadata=metadata, + required=required, + shape=shape, + ) diff --git a/pandera/api/tensordict/model_config.py b/pandera/api/tensordict/model_config.py new file mode 100644 index 000000000..638f2d138 --- /dev/null +++ b/pandera/api/tensordict/model_config.py @@ -0,0 +1,14 @@ +"""Class-based TensorDict model API configuration.""" + +from typing import Any + +from pandera.api.dataframe.model_config import BaseConfig as _BaseConfig + + +class BaseConfig(_BaseConfig): + """Define TensorDictSchema-wide options. + + *new in 0.19.0* + """ + + batch_size: tuple[int, ...] | None = None diff --git a/pandera/api/tensordict/types.py b/pandera/api/tensordict/types.py new file mode 100644 index 000000000..f14a5b06b --- /dev/null +++ b/pandera/api/tensordict/types.py @@ -0,0 +1,25 @@ +"""Type definitions for pandera tensordict integration.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Union + +from pandera.api.checks import Check + +try: + import torch +except ImportError: + torch = None + +if torch is not None: + from tensordict import TensorDict, tensorclass + + +def is_tensordict(obj: Any) -> bool: + """Check if object is a TensorDict or tensorclass.""" + try: + from tensordict import TensorDict, tensorclass + + return isinstance(obj, (TensorDict, tensorclass)) + except ImportError: + return False diff --git a/pandera/backends/pandas/container.py b/pandera/backends/pandas/container.py index 12d03895b..5c8a2b712 100644 --- a/pandera/backends/pandas/container.py +++ b/pandera/backends/pandas/container.py @@ -16,9 +16,7 @@ PandasSchemaBackend, _parsed_column_values, ) -from pandera.backends.pandas.error_formatters import ( - reshape_failure_cases, -) +from pandera.backends.pandas.error_formatters import reshape_failure_cases from pandera.backends.utils import convert_uniquesettings from pandera.config import ValidationScope from pandera.engines import pandas_engine diff --git a/pandera/backends/pyspark/base.py b/pandera/backends/pyspark/base.py index 5e809a417..56ff06517 100644 --- a/pandera/backends/pyspark/base.py +++ b/pandera/backends/pyspark/base.py @@ -2,13 +2,7 @@ import warnings from collections.abc import Iterable -from typing import ( - Any, - NamedTuple, - Optional, - TypeVar, - Union, -) +from typing import Any, NamedTuple, Optional, TypeVar, Union from pyspark.sql import DataFrame from pyspark.sql.functions import col diff --git a/pandera/backends/tensordict/__init__.py b/pandera/backends/tensordict/__init__.py new file mode 100644 index 000000000..ac84da606 --- /dev/null +++ b/pandera/backends/tensordict/__init__.py @@ -0,0 +1 @@ +"""Pandera TensorDict backend.""" diff --git a/pandera/backends/tensordict/base.py b/pandera/backends/tensordict/base.py new file mode 100644 index 000000000..6c37ef3e9 --- /dev/null +++ b/pandera/backends/tensordict/base.py @@ -0,0 +1,399 @@ +"""TensorDict parsing, validation, and error-reporting backends.""" + +import copy +from collections import defaultdict +from typing import Any + +from pandera import errors +from pandera.api.base.error_handler import ErrorHandler, get_error_category +from pandera.backends.base import BaseSchemaBackend, CoreCheckResult +from pandera.errors import SchemaError, SchemaErrorReason + +try: + import torch +except ImportError: + torch = None + + +def _is_tensordict(obj: Any) -> bool: + """Check if object is a TensorDict or tensorclass.""" + if torch is None: + return False + try: + from tensordict import TensorDict + + # Check for TensorDict first (more common case) + if isinstance(obj, TensorDict): + return True + + # Check for tensorclass - use _is_tensorclass attribute which is set + # on all tensorclass instances + if hasattr(obj, "_is_tensorclass") and obj._is_tensorclass: + return True + + return False + except ImportError: + return False + + +def _get_tensor(obj: Any, key: str) -> Any: + """Get tensor from TensorDict or tensorclass object.""" + # For TensorDict, use __getitem__ + try: + return obj[key] + except (KeyError, TypeError, ValueError): + pass + + # For tensorclass, use attribute access + if hasattr(obj, key): + return getattr(obj, key) + + raise KeyError(f"Key '{key}' not found in object") + + +class TensorDictSchemaBackend(BaseSchemaBackend): + """Backend for TensorDict schema validation.""" + + def preprocess(self, check_obj, inplace: bool = False): + """Preprocesses a check object before applying check functions.""" + if not inplace: + if _is_tensordict(check_obj): + check_obj = check_obj.clone() + else: + check_obj = check_obj.copy() + return check_obj + + def validate( + self, + check_obj, + schema, + *, + head: int | None = None, + tail: int | None = None, + sample: int | None = None, + random_state: int | None = None, + lazy: bool = False, + inplace: bool = False, + ): + """Validate a TensorDict against the schema.""" + error_handler = ErrorHandler(lazy) + + check_obj = self.preprocess(check_obj, inplace=inplace) + + error_handler = self.run_checks_and_handle_errors( + error_handler, schema, check_obj, lazy + ) + + if error_handler.collected_errors: + raise errors.SchemaErrors( + schema=schema, + schema_errors=error_handler.schema_errors, + data=check_obj, + ) + + return check_obj + + def coerce_dtype(self, check_obj, schema=None): + """Coerce the dtype of tensors in the TensorDict.""" + return check_obj + + def run_checks_and_handle_errors( + self, error_handler, schema, check_obj, lazy + ): + """Run all checks and collect errors.""" + self._check_batch_size(check_obj, schema, error_handler, lazy) + self._check_keys(check_obj, schema, error_handler, lazy) + self._check_dtypes_and_shapes(check_obj, schema, error_handler, lazy) + self._run_value_checks(check_obj, schema, error_handler, lazy) + + return error_handler + + def _check_batch_size(self, check_obj, schema, error_handler, lazy): + """Validate batch_size matches.""" + if schema.batch_size is None: + return + + actual_batch_size = check_obj.batch_size + expected_batch_size = schema.batch_size + + if len(actual_batch_size) != len(expected_batch_size): + error_msg = ( + f"Expected batch_size {expected_batch_size}, " + f"got {actual_batch_size}" + ) + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_DATATYPE, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + return + + for i, (exp, act) in enumerate( + zip(expected_batch_size, actual_batch_size) + ): + if exp is not None and exp != act: + error_msg = f"Expected batch_size[{i}]={exp}, got batch_size[{i}]={act}" + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_DATATYPE, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + + def _check_keys(self, check_obj, schema, error_handler, lazy): + """Validate required keys exist.""" + for key in schema.keys: + # Check if key exists - use different methods for TensorDict vs tensorclass + is_key_present = False + + try: + from tensordict import TensorDict + + # For TensorDict, use 'in' operator + if isinstance(check_obj, TensorDict): + is_key_present = key in check_obj + else: + # For tensorclass, use keys() method or attribute access + if hasattr(check_obj, "keys"): + is_key_present = key in check_obj.keys() + elif hasattr(check_obj, key): + is_key_present = True + except ImportError: + pass + + if not is_key_present: + error_msg = f"Missing key '{key}' in TensorDict" + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.COLUMN_NOT_IN_DATAFRAME, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + + def _check_dtypes_and_shapes(self, check_obj, schema, error_handler, lazy): + """Validate tensor dtypes and shapes.""" + for key, tensor_schema in schema.keys.items(): + try: + tensor = _get_tensor(check_obj, key) + except KeyError: + # Key doesn't exist (should already be caught by _check_keys) + continue + + if not isinstance(tensor, torch.Tensor): + error_msg = f"Key '{key}' is not a torch.Tensor" + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_DATATYPE, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + continue + + if tensor_schema.dtype is not None: + expected_dtype = tensor_schema.dtype.type + actual_dtype = tensor.dtype + if actual_dtype != expected_dtype: + error_msg = ( + f"Key '{key}': expected dtype {expected_dtype}, " + f"got {actual_dtype}" + ) + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_DATATYPE, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + + if tensor_schema.shape is not None: + actual_shape = tuple(tensor.shape) + expected_shape = tensor_schema.shape + + if len(actual_shape) != len(expected_shape): + error_msg = ( + f"Key '{key}': expected shape {expected_shape}, " + f"got {actual_shape}" + ) + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_DATATYPE, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + continue + + for i, (exp, act) in enumerate( + zip(expected_shape, actual_shape) + ): + if exp is not None and exp != act: + error_msg = ( + f"Key '{key}': expected shape[{i}]={exp}, " + f"got shape[{i}]={act}" + ) + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_DATATYPE, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + + def _run_value_checks(self, check_obj, schema, error_handler, lazy): + """Run value checks on tensor values.""" + from pandera.api.base.checks import CheckResult + + for key, tensor_schema in schema.keys.items(): + try: + tensor = _get_tensor(check_obj, key) + except KeyError: + continue + + if not tensor_schema.checks: + continue + + # For value checks, we need to pass the tensor data + tensor_data = {key: tensor} + + for check_index, check in enumerate(tensor_schema.checks): + try: + # Use TensorDictCheckBackend directly + from pandera.backends.tensordict.checks import ( + TensorDictCheckBackend, + ) + + check_backend = TensorDictCheckBackend(check) + + # Apply check to the tensor (not dict) + check_result = check_backend.apply(tensor) + + # Postprocess result + check_result = check_backend.postprocess( + tensor, check_result + ) + except Exception as exc: + check_result = CoreCheckResult( + passed=False, + check=check, + check_index=check_index, + reason_code=SchemaErrorReason.CHECK_ERROR, + message=str(exc), + ) + + if not check_result.check_passed: + error_msg = ( + f"Check '{check}' failed for key '{key}': check failed" + ) + error = SchemaError( + schema=schema, + data=check_obj, + message=error_msg, + check=check, + check_index=check_index, + reason_code=SchemaErrorReason.CHECK_ERROR, + ) + error_handler.collect_error( + get_error_category(error.reason_code), + error.reason_code, + error, + ) + + def run_check(self, check_obj, schema, check, check_index, *args): + """Run a single check on the check object.""" + raise NotImplementedError("Use _run_value_checks instead") + + def run_checks(self, check_obj, schema): + """Run a list of checks on the check object.""" + raise NotImplementedError("Use _run_value_checks instead") + + def run_schema_component_checks( + self, check_obj, schema, schema_components, lazy + ): + """Run checks for all schema components.""" + raise NotImplementedError("Use _run_value_checks instead") + + def check_name(self, check_obj, schema): + """Core check that checks the name of the check object.""" + raise NotImplementedError + + def check_nullable(self, check_obj, schema): + """Core check that checks the nullability of a check object.""" + raise NotImplementedError + + def check_unique(self, check_obj, schema): + """Core check that checks the uniqueness of values in a check object.""" + raise NotImplementedError + + def check_dtype(self, check_obj, schema): + """Core check that checks the data type of a check object.""" + raise NotImplementedError + + def failure_cases_metadata( + self, schema_name: str, schema_errors: list[SchemaError] + ): + """Get failure cases metadata for lazy validation.""" + from collections import defaultdict + + from pandera.api.base.error_handler import ErrorHandler + from pandera.errors import FailureCaseMetadata + + error_dicts = {} + error_handler = ErrorHandler() + error_handler.collect_errors(schema_errors) + + if error_handler.collected_errors: + error_dicts = error_handler.summarize(schema_name=schema_name) + error_dicts = dict(error_dicts) + + failure_cases = [ + err.failure_cases for err in schema_errors if err.failure_cases + ] + + error_counts: dict[str, int] = defaultdict(int) + for error in error_handler.collected_errors: + error_counts[error["reason_code"].name] += 1 + + return FailureCaseMetadata( + failure_cases=failure_cases, + message=error_dicts, + error_counts=dict(error_counts), + ) + + def drop_invalid_rows(self, check_obj, error_handler): + """Remove invalid elements in a check_obj.""" + raise NotImplementedError( + "drop_invalid_rows is not applicable for TensorDict" + ) diff --git a/pandera/backends/tensordict/builtin_checks.py b/pandera/backends/tensordict/builtin_checks.py new file mode 100644 index 000000000..c89f55506 --- /dev/null +++ b/pandera/backends/tensordict/builtin_checks.py @@ -0,0 +1,135 @@ +"""Built-in checks for TensorDict.""" + +import sys + +# Import torch first to get the actual type, not string annotation +try: + import torch + + TORCH_AVAILABLE = True +except ImportError: + TORCH_AVAILABLE = False + torch = None + +from collections.abc import Callable, Iterable +from functools import partial +from typing import Any + +try: + from pandera.api.extensions import register_builtin_check + + if TORCH_AVAILABLE: + + @register_builtin_check( + aliases=["gt"], + error="greater_than({min_value})", + ) + def greater_than(data: torch.Tensor, min_value: Any) -> torch.Tensor: + """Ensure values are strictly greater than a minimum value. + + :param data: Input tensor data. + :param min_value: Lower bound to be exceeded. + """ + return data > min_value + + @register_builtin_check( + aliases=["ge"], + error="greater_than_or_equal_to({min_value})", + ) + def greater_than_or_equal_to( + data: torch.Tensor, min_value: Any + ) -> torch.Tensor: + """Ensure all values are greater than or equal to a minimum value. + + :param data: Input tensor data. + :param min_value: Allowed minimum value. + """ + return data >= min_value + + @register_builtin_check( + aliases=["lt"], + error="less_than({max_value})", + ) + def less_than(data: torch.Tensor, max_value: Any) -> torch.Tensor: + """Ensure all values are strictly less than a maximum value. + + :param data: Input tensor data. + :param max_value: Allowed maximum value. + """ + return data < max_value + + @register_builtin_check( + aliases=["le"], + error="less_than_or_equal_to({max_value})", + ) + def less_than_or_equal_to( + data: torch.Tensor, max_value: Any + ) -> torch.Tensor: + """Ensure all values are less than or equal to a maximum value. + + :param data: Input tensor data. + :param max_value: Allowed maximum value. + """ + return data <= max_value + + @register_builtin_check( + aliases=["in_range"], + error="in_range({min_value}, {max_value})", + ) + def in_range( + data: torch.Tensor, + min_value: Any, + max_value: Any, + include_min: bool = True, + include_max: bool = True, + ) -> torch.Tensor: + """Ensure all values are within a range. + + :param data: Input tensor data. + :param min_value: Minimum allowed value. + :param max_value: Maximum allowed value. + :param include_min: If True, min_value is included in the range. + :param include_max: If True, max_value is included in the range. + """ + if include_min and include_max: + return (data >= min_value) & (data <= max_value) + elif include_min: + return (data >= min_value) & (data < max_value) + elif include_max: + return (data > min_value) & (data <= max_value) + else: + return (data > min_value) & (data < max_value) + + @register_builtin_check( + aliases=["isin"], + error="isin({allowed_values})", + ) + def isin(data: torch.Tensor, allowed_values: Iterable) -> torch.Tensor: + """Ensure all values are in a set of allowed values. + + :param data: Input tensor data. + :param allowed_values: Set of allowed values. + """ + return torch.isin(data, torch.as_tensor(list(allowed_values))) + +except ImportError: + # extensions not available yet + def register_builtin_check( + fn: Any = None, + strategy: Callable | None = None, + _check_cls: type = None, # type: ignore[assignment] + aliases: list[str] | None = None, + **kwargs: Any, + ) -> Any: + def decorator(f): + return f + + if fn is None: + return partial( + register_builtin_check, + strategy=strategy, + _check_cls=_check_cls, + aliases=aliases, + **kwargs, + ) + return decorator(fn) diff --git a/pandera/backends/tensordict/checks.py b/pandera/backends/tensordict/checks.py new file mode 100644 index 000000000..2296020a5 --- /dev/null +++ b/pandera/backends/tensordict/checks.py @@ -0,0 +1,119 @@ +"""Check backend for TensorDict.""" + +from functools import partial +from typing import Any + +from pandera.api.base.checks import CheckResult +from pandera.api.checks import Check +from pandera.backends.base import BaseCheckBackend, CoreCheckResult +from pandera.errors import SchemaErrorReason + + +class TensorDictCheckBackend(BaseCheckBackend): + """Check backend for TensorDict.""" + + def __init__(self, check: Check): + """Initializes a check backend object.""" + super().__init__(check) + self.check = check + self.check_fn = ( + partial(check._check_fn, **check._check_kwargs) + if check._check_fn is not None + else None + ) + + def preprocess(self, check_obj: Any, key: Any = None) -> Any: + """Preprocesses the check object before applying the check function.""" + # For tensor data, no preprocessing needed + return check_obj + + def apply(self, check_obj: Any) -> CheckResult: + """Apply the check function to the check object. + + :param check_obj: Tensor data to validate. + :returns: CheckResult with passed status and output. + """ + check_obj = self.preprocess(check_obj) + + if self.check_fn is None: + return CheckResult( + check_output=True, + check_passed=True, + checked_object=check_obj, + failure_cases=None, + ) + + # Apply the check function + try: + result = self.check_fn(check_obj) + except Exception as exc: + return CheckResult( + check_output=False, + check_passed=False, + checked_object=check_obj, + failure_cases=None, + ) + + # Reduce to boolean + passed = self._reduce_to_bool(result) + + if not passed: + failure_cases = self._get_failure_cases(check_obj, result) + else: + failure_cases = None + + return CheckResult( + check_output=result, + check_passed=passed, + checked_object=check_obj, + failure_cases=failure_cases, + ) + + def _reduce_to_bool(self, result: Any) -> bool: + """Reduce a boolean tensor/array to a single boolean value.""" + try: + import torch + + if isinstance(result, torch.Tensor): + return result.all().item() + except ImportError: + pass + + try: + import numpy as np + + if isinstance(result, np.ndarray): + return result.all().item() + except ImportError: + pass + + # Assume it's already a boolean + return bool(result) + + def _get_failure_cases(self, check_obj: Any, result: Any) -> Any | None: + """Get the failure cases from the check result.""" + try: + import torch + + if isinstance(check_obj, torch.Tensor): + # Get indices where check failed + if isinstance(result, torch.Tensor): + failure_mask = ~result + return check_obj[failure_mask] + except ImportError: + pass + + return None + + def postprocess( + self, check_obj: Any, check_output: CheckResult + ) -> CheckResult: + """Postprocesses the result of applying the check function.""" + if isinstance(check_output, CheckResult): + return check_output + return CheckResult( + check_output=check_output, + check_passed=bool(check_output), + checked_object=check_obj, + failure_cases=None, + ) diff --git a/pandera/backends/tensordict/register.py b/pandera/backends/tensordict/register.py new file mode 100644 index 000000000..a59aa623c --- /dev/null +++ b/pandera/backends/tensordict/register.py @@ -0,0 +1,39 @@ +"""Register TensorDict backends.""" + +from functools import lru_cache + + +@lru_cache +def register_tensordict_backends() -> None: + """Register TensorDict backends. + + This function is called at schema initialization in the _register_*_backends + method to register validation backends for TensorDict and tensorclass objects. + """ + from pandera.api.checks import Check + from pandera.api.tensordict.container import TensorDictSchema + + # Import builtin checks to register check functions + from pandera.backends.tensordict import builtin_checks + from pandera.backends.tensordict.base import ( + TensorDictSchemaBackend, + _is_tensordict, + ) + from pandera.backends.tensordict.checks import TensorDictCheckBackend + + try: + import torch + from tensordict import TensorDict + + # Register schema backend for TensorDict (has a fixed type) + TensorDictSchema.register_backend(TensorDict, TensorDictSchemaBackend) + + # For tensorclass, we handle it in get_backend() since each decorated + # class creates a unique type. We use TensorDict's backend since they + # share the same interface. + + # Register check backend for tensors + Check.register_backend(torch.Tensor, TensorDictCheckBackend) + + except ImportError: + pass diff --git a/pandera/backends/xarray/data_tree.py b/pandera/backends/xarray/data_tree.py index bf5e35584..83c53693d 100644 --- a/pandera/backends/xarray/data_tree.py +++ b/pandera/backends/xarray/data_tree.py @@ -5,10 +5,7 @@ from typing import Any from pandera.api.base.error_handler import ErrorHandler, get_error_category -from pandera.api.xarray.container import ( - DatasetSchema, - DataTreeSchema, -) +from pandera.api.xarray.container import DatasetSchema, DataTreeSchema from pandera.backends.base import CoreCheckResult from pandera.backends.xarray.base import XarraySchemaBackend from pandera.backends.xarray.container import DatasetSchemaBackend diff --git a/pandera/dtypes.py b/pandera/dtypes.py index 762fd96a4..d50dea116 100644 --- a/pandera/dtypes.py +++ b/pandera/dtypes.py @@ -7,13 +7,7 @@ import inspect from abc import ABC from collections.abc import Callable, Iterable -from typing import ( - Any, - Literal, - Optional, - TypeVar, - Union, -) +from typing import Any, Literal, Optional, TypeVar, Union from typing_extensions import overload diff --git a/pandera/engines/pandas_engine.py b/pandera/engines/pandas_engine.py index a7f47aa1c..5f6950769 100644 --- a/pandera/engines/pandas_engine.py +++ b/pandera/engines/pandas_engine.py @@ -12,14 +12,7 @@ import sys import warnings from collections.abc import Callable, Iterable -from typing import ( - Any, - Literal, - NamedTuple, - Optional, - Union, - cast, -) +from typing import Any, Literal, NamedTuple, Optional, Union, cast import numpy as np import pandas as pd diff --git a/pandera/engines/polars_engine.py b/pandera/engines/polars_engine.py index 13f6d6baa..ad5011ee9 100644 --- a/pandera/engines/polars_engine.py +++ b/pandera/engines/polars_engine.py @@ -8,14 +8,7 @@ import logging import warnings from collections.abc import Iterable, Mapping, Sequence -from typing import ( - Any, - Literal, - Optional, - TypedDict, - Union, - overload, -) +from typing import Any, Literal, Optional, TypedDict, Union, overload import polars as pl from packaging import version diff --git a/pandera/engines/tensordict_engine.py b/pandera/engines/tensordict_engine.py new file mode 100644 index 000000000..263fb7a84 --- /dev/null +++ b/pandera/engines/tensordict_engine.py @@ -0,0 +1,150 @@ +"""TensorDict engine and data types.""" + +import dataclasses +from typing import Any, Union + +from pandera import dtypes +from pandera.dtypes import immutable +from pandera.engines import engine + +try: + import torch +except ImportError: + torch = None + + +if torch is not None: + + @immutable(init=True) + class DataType(dtypes.DataType): + """DataType for boxing PyTorch tensor dtypes.""" + + type: torch.dtype = dataclasses.field( + default=torch.float32, repr=False, init=False + ) + + def __init__(self, dtype: Any): + super().__init__() + if isinstance(dtype, torch.dtype): + actual_dtype = dtype + elif isinstance(dtype, str): + dtype_str = dtype.replace("torch.", "") + actual_dtype = getattr(torch, dtype_str, None) + if actual_dtype is None: + raise ValueError(f"Unknown torch dtype: {dtype}") + else: + actual_dtype = torch.dtype(dtype) + object.__setattr__(self, "type", actual_dtype) + + def __post_init__(self): + if isinstance(self.type, torch.dtype): + return + actual_dtype = getattr(torch, str(self.type), None) + if actual_dtype is not None: + object.__setattr__(self, "type", actual_dtype) + + def coerce(self, data_container: torch.Tensor) -> torch.Tensor: + """Coerce tensor to the specified dtype.""" + return data_container.type(self.type) + + def coerce_value(self, value: Any) -> Any: + """Coerce a value to the particular type.""" + return torch.tensor(value, dtype=self.type) + + def try_coerce(self, data_container: torch.Tensor) -> torch.Tensor: + try: + return self.coerce(data_container) + except Exception as exc: + from pandera import errors + + raise errors.ParserError( + f"Could not coerce tensor to type {self.type}", + failure_cases=None, + ) from exc + + def __str__(self) -> str: + return str(self.type) + + def __repr__(self) -> str: + return f"DataType({self})" + + class Engine(metaclass=engine.Engine, base_pandera_dtypes=DataType): + """PyTorch TensorDict data type engine.""" + + @classmethod + def dtype(cls, data_type: Any) -> dtypes.DataType: + """Convert input into a PyTorch-compatible Pandera DataType.""" + if isinstance(data_type, type) and issubclass(data_type, DataType): + return DataType(str(data_type.type)) + try: + return engine.Engine.dtype(cls, data_type) + except (TypeError, ValueError): + try: + if isinstance(data_type, torch.dtype): + return DataType(str(data_type)) + elif isinstance(data_type, str): + return DataType(data_type) + else: + return DataType(data_type) + except (TypeError, ValueError): + raise ValueError( + f"Data type '{data_type}' not understood for TensorDict. " + f"Expected a torch.dtype or string like 'float32'." + ) from None + + @classmethod + def register_dtype(cls, pandera_dtype_cls: type[DataType]): + """Register a Pandera DataType for PyTorch dtypes.""" + cls._check_source_dtype(pandera_dtype_cls) + equivalents = {} + strict_equivalents: dict[Any, DataType] = {} + for source_dtype in (pandera_dtype_cls.type,): + if source_dtype is not None: + equivalents[source_dtype] = pandera_dtype_cls # type: ignore[assignment] + + if equivalents: + registry = cls._registry[cls] + registry.equivalents.update(equivalents) + registry.strict_equivalents.update(strict_equivalents) + + def _register_torch_dtypes(): + """Register all torch dtypes.""" + torch_dtype_names = [ + "float32", + "float64", + "int32", + "int64", + "int32", + "int16", + "int8", + "uint8", + "bool", + "bfloat16", + "complex64", + "complex128", + "float16", + "quint8", + "qint8", + "qint32", + ] + + for dtype_name in torch_dtype_names: + try: + torch_dtype = getattr(torch, dtype_name, None) + if torch_dtype is not None: + Engine.register_dtype(DataType(dtype_name)) + except (AttributeError, ValueError): + pass + + _register_torch_dtypes() + +else: + + class _DataTypePlaceholder: + pass + + class _EnginePlaceholder: + pass + + DataType = None # type: ignore[misc, assignment] + Engine = None # type: ignore[misc, assignment] diff --git a/pandera/io/xarray_io.py b/pandera/io/xarray_io.py index 730899d1d..78ac916d5 100644 --- a/pandera/io/xarray_io.py +++ b/pandera/io/xarray_io.py @@ -239,9 +239,7 @@ def serialize_dataset_schema( :returns: dict representation of the schema. """ from pandera import __version__ - from pandera.schema_statistics.xarray import ( - get_dataset_schema_statistics, - ) + from pandera.schema_statistics.xarray import get_dataset_schema_statistics stats = get_dataset_schema_statistics(dataset_schema) diff --git a/pandera/schema_inference/geopandas.py b/pandera/schema_inference/geopandas.py index 425943f5d..9fdf7ff56 100644 --- a/pandera/schema_inference/geopandas.py +++ b/pandera/schema_inference/geopandas.py @@ -11,9 +11,7 @@ from pandera.schema_inference.pandas import ( infer_dataframe_schema as _infer_pandas_df_schema, ) -from pandera.schema_inference.pandas import ( - infer_series_schema, -) +from pandera.schema_inference.pandas import infer_series_schema def infer_geodataframe_schema(df: pd.DataFrame) -> GeoDataFrameSchema: diff --git a/pandera/strategies/pandas_strategies.py b/pandera/strategies/pandas_strategies.py index 9d0c0d5c4..7a04d9da8 100644 --- a/pandera/strategies/pandas_strategies.py +++ b/pandera/strategies/pandas_strategies.py @@ -17,13 +17,7 @@ from collections.abc import Callable, Sequence from copy import deepcopy from functools import partial, wraps -from typing import ( - Any, - Optional, - TypeVar, - Union, - cast, -) +from typing import Any, Optional, TypeVar, Union, cast import numpy as np import pandas as pd diff --git a/pandera/tensordict.py b/pandera/tensordict.py new file mode 100644 index 000000000..65380b0b5 --- /dev/null +++ b/pandera/tensordict.py @@ -0,0 +1,36 @@ +"""Pandera TensorDict API.""" + +from __future__ import annotations + +# mypy: disable-error-code=attr-defined +from typing import TYPE_CHECKING + +from pandera import errors + +try: + from pandera.engines.tensordict_engine import DataType as _DataType + + DataType: type[_DataType] | None = _DataType # type: ignore[misc] +except ImportError: + DataType = None # type: ignore[misc, assignment] + +# Import actual implementations (not just for type checking) +from pandera.api.tensordict.components import Tensor +from pandera.api.tensordict.container import TensorDictSchema +from pandera.api.tensordict.model import TensorDictModel +from pandera.api.tensordict.model_components import Field + +__all__ = [ + "Tensor", + "TensorDictSchema", + "TensorDictModel", + "Field", + "errors", + "DataType", + # Error classes + "SchemaError", + "SchemaErrors", +] + +# Import error classes for convenience +from pandera.errors import SchemaError, SchemaErrors diff --git a/pandera/typing/__init__.py b/pandera/typing/__init__.py index d27ad18bc..9459278a4 100644 --- a/pandera/typing/__init__.py +++ b/pandera/typing/__init__.py @@ -50,13 +50,7 @@ @lru_cache def get_dataframe_types(): - from pandera.typing import ( - dask, - geopandas, - modin, - pyspark, - pyspark_sql, - ) + from pandera.typing import dask, geopandas, modin, pyspark, pyspark_sql dataframe_types: set[type] = {DataFrame} if dask.DASK_INSTALLED: @@ -79,12 +73,7 @@ def get_dataframe_types(): @lru_cache def get_series_types(): - from pandera.typing import ( - dask, - geopandas, - modin, - pyspark, - ) + from pandera.typing import dask, geopandas, modin, pyspark series_types: set[type] = {Series} if dask.DASK_INSTALLED: diff --git a/pandera/typing/pandas.py b/pandera/typing/pandas.py index 4f5875d07..1f53e85cb 100644 --- a/pandera/typing/pandas.py +++ b/pandera/typing/pandas.py @@ -25,9 +25,7 @@ IndexBase, SeriesBase, ) -from pandera.typing.common import ( - GenericDtype as _CommonGenericDtype, -) +from pandera.typing.common import GenericDtype as _CommonGenericDtype from pandera.typing.formats import Formats try: diff --git a/pyproject.toml b/pyproject.toml index 0ced750a6..a7b42c1ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,6 +98,10 @@ xarray = [ "xarray >= 2024.10.0", "numpy >= 1.24.4", ] +torch = [ + "torch", + "tensordict", +] all = [ "hypothesis >= 6.92.7", "scipy", @@ -119,6 +123,8 @@ all = [ "polars >= 0.20.0", "xarray >= 2024.10.0", "numpy >= 1.24.4", + "torch", + "tensordict", ] [dependency-groups] @@ -153,6 +159,7 @@ testing = [ ] docs = [ "setuptools", + "duckdb", "pyarrow-hotfix", "sphinx<9", "sphinx-design", @@ -164,11 +171,14 @@ docs = [ "sphinx-docsearch", "grpcio", "ray", + "torch", + "tensordict", "types-click", "types-pytz", "types-pyyaml", "types-requests", "types-setuptools", + "xdoctest", ] [tool.setuptools] diff --git a/requirements.txt b/requirements.txt index 72d4f94b0..7189fc2fd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,6 +17,8 @@ scipy-stubs pyspark[connect] >= 3.2.0 polars >= 0.20.0 xarray >= 2024.10.0 +torch +tensordict modin protobuf geopandas diff --git a/specs/pytorch-tensordict.md b/specs/pytorch-tensordict.md index 8f8001ebb..28b8f9c33 100644 --- a/specs/pytorch-tensordict.md +++ b/specs/pytorch-tensordict.md @@ -49,20 +49,20 @@ def validate_batch(td: TensorDict): """Manual validation - error-prone, hard to maintain.""" # Check batch size assert td.batch_size[0] == 32, f"Expected batch_size[0]=32, got {td.batch_size[0]}" - + # Check keys exist assert "observation" in td, "Missing 'observation' key" assert "action" in td, "Missing 'action' key" - + # Check dtypes assert td["observation"].dtype == torch.float32, f"Expected float32, got {td['observation'].dtype}" assert td["action"].dtype == torch.float32, f"Expected float32, got {td['action'].dtype}" - + # Check value ranges (normalized observations) obs = td["observation"] assert obs.min() >= -1.0, f"observation min value {obs.min()} < -1.0" assert obs.max() <= 1.0, f"observation max value {obs.max()} > 1.0" - + # Check action bounds action = td["action"] assert action.min() >= -2.0, f"action min value {action.min()} < -2.0" @@ -158,38 +158,32 @@ except pa.SchemaErrors as e: #### Example 3: Declarative API with `TensorDictModel` -For a more Pydantic-like experience, use class-based models: +For a more Pydantic-like experience, use class-based models. **Type annotations specify the dtype directly**, and `pa.Field` defines additional constraints: ```python import torch import pandera.tensordict as pa -from pandera import Field, Check +from pandera import Check class PolicyModel(pa.TensorDictModel): """Schema for a policy network output.""" - - logits: torch.Tensor = Field( - dtype=torch.float32, + + # Dtype is specified in the type annotation, Field defines constraints + logits: torch.float32 = pa.Field( shape=(None, 10), - checks=Check.in_range(-5.0, 5.0) + ge=-5.0, + le=5.0 ) - value: torch.Tensor = Field( - dtype=torch.float32, + value: torch.float32 = pa.Field( shape=(None,), - checks=Check.greater_than(0.0) + gt=0.0 ) - + class Config: batch_size = (64,) -# Validate using the model -policy_model = PolicyModel() -batch = TensorDict({ - "logits": torch.randn(64, 10), - "value": torch.randn(64).abs() # Ensure positive -}, batch_size=[64]) - -validated = policy_model.validate(batch) +# Validate using the model - schema is built automatically from class definition +validated_batch = PolicyModel.validate(batch) ``` #### Example 4: Validating `tensorclass` Objects @@ -312,10 +306,10 @@ def save_dataset(td: TensorDict, path: str): def load_and_validate_dataset(path: str) -> TensorDict: """Load dataset and validate before use.""" td = TensorDict.load(path) - + # Validate schema compliance validated = dataset_schema.validate(td) - + return validated # Create and save a large dataset (e.g., ImageNet subset) @@ -387,7 +381,7 @@ def save_trajectory(tc: TrajectoryData, path: str): def load_trajectory(path: str) -> TrajectoryData: """Load and validate tensorclass from disk.""" tc = TrajectoryData.load(path) - + # Validate - catches any corruption from disk I/O return trajectory_schema.validate(tc) @@ -437,15 +431,15 @@ def validate_dataset_directory(data_dir: Path) -> list[str]: """Validate all TensorDict files in a directory.""" errors = [] schema = large_dataset_schema - + for td_file in sorted(data_dir.glob("*.pt")): td = TensorDict.load(td_file) - + try: schema.validate(td, lazy=True) except pa.SchemaErrors as e: errors.append(f"{td_file}: {e}") - + return errors # Check entire dataset before starting training @@ -538,10 +532,10 @@ replay_buffer = ReplayBuffer( # Validate sampled data before training def train_step(): batch = replay_buffer.sample() - + # Validate the batch - catch data corruption or storage errors validated_batch = replay_schema.validate(batch) - + # Now safe to pass to loss computation loss = compute_ppo_loss(validated_batch) loss.backward() @@ -560,25 +554,25 @@ from pandera import Check # Schema for actor-critic network output class ActorCriticOutput(pa.TensorDictModel): """Output from an actor-critic network.""" - + action_mean: torch.Tensor = pa.Field( dtype=torch.float32, shape=(None, 4), # 4D action space - checks=Check.greater_than_or_equal_to(-2.0), - Check.less_than_or_equal_to(2.0), + ge=-2.0, + le=2.0, ) action_std: torch.Tensor = pa.Field( dtype=torch.float32, shape=(None, 4), - checks=Check.greater_than(1e-6), # Standard deviation > 0 + gt=1e-6, # Standard deviation > 0 ) value: torch.Tensor = pa.Field( dtype=torch.float32, shape=(None,), - checks=Check.greater_than(-100.0), # Reasonable value bounds - Check.less_than(100.0), + gt=-100.0, # Reasonable value bounds + lt=100.0, ) - + class Config: batch_size = (32,) @@ -614,7 +608,7 @@ training_schema = pa.TensorDictSchema( "action": pa.Tensor(dtype=torch.float32, shape=(None, 6)), "reward": pa.Tensor(dtype=torch.float32, shape=(None,)), "done": pa.Tensor(dtype=torch.bool, shape=(None,)), - + # PPO-specific data "log_prob": pa.Tensor(dtype=torch.float32, shape=(None,)), "value": pa.Tensor(dtype=torch.float32, shape=(None,)), @@ -651,19 +645,19 @@ trainer_validator = pa.TensorDictSchema( def distributed_train_step(collector_rank, trainer_rank): # Collect experiences experience = collector.collect_experiences() - + # Validate before sending to replay buffer validated_exp = collector_validator.validate(experience) - + # Send to buffer (RRef) replay_buffer_ref.add(validated_exp) - + # Trainer pulls batch batch = replay_buffer_ref.sample() - + # Validate before training validated_batch = trainer_validator.validate(batch) - + # Train optimizer.zero_grad() loss = compute_ppo_loss(validated_batch) @@ -840,6 +834,127 @@ class MyTensorDictModel(pa.TensorDictModel): batch_size = (None,) ``` + + +### 4.5 Class-Based Model Implementation + +The `TensorDictModel` class follows the same pattern as `xarray.DatasetModel` and `pandas.DataFrameModel`. The API supports direct usage of `pa.Field` in type annotations without requiring custom `_field` classmethods or other complex patterns. + +**Key principles:** +- Use PyTorch dtypes (e.g., `torch.float32`, `torch.int64`) in type annotations +- Define constraints using `pa.Field()` descriptor +- Configuration via nested `Config` class + +#### Example: Complete Model Definition + +```python +import torch +import pandera.tensordict as pa +from pandera import Check + +class RLModel(pa.TensorDictModel): + """Reinforcement learning model output schema.""" + + # Dtype in annotation, constraints in Field + observation: torch.float32 = pa.Field(shape=(None, 128)) + action: torch.int64 = pa.Field( + shape=(None,), + checks=Check.isin([0, 1, 2, 3]) + ) + reward: torch.float32 = pa.Field() + + class Config: + batch_size = (32,) + +# Access schema directly +schema = RLModel.to_schema() + +# Or validate directly +batch = TensorDict({ + "observation": torch.randn(32, 128), + "action": torch.randint(0, 4, (32,)), + "reward": torch.randn(32) +}, batch_size=[32]) + +validated = RLModel.validate(batch) +``` + +#### Example: Class-Based Model with Field-Level Checks + +```python +class ActorCriticOutput(pa.TensorDictModel): + """Actor-critic network output.""" + + action_mean: torch.float32 = pa.Field( + shape=(None, 4), + ge=-2.0, + le=2.0, + ) + action_std: torch.float32 = pa.Field( + shape=(None, 4), + gt=1e-6, + ) + value: torch.float32 = pa.Field( + shape=(None,), + gt=-100.0, + lt=100.0, + ) + + class Config: + batch_size = (32,) +``` + +#### Example: Lazy Validation with Class-Based Model + +```python +# Create invalid data +invalid_batch = TensorDict({ + "observation": torch.randn(32, 128), + "action": torch.randint(0, 4, (16,)), # Wrong batch size! + "reward": torch.randn(32) +}, batch_size=[32]) + +try: + RLModel.validate(invalid_batch, lazy=True) +except pa.SchemaErrors as e: + print(e) # Shows all validation errors +``` + +#### Example: TensorDictSchema vs TensorDictModel + +For simple use cases, `TensorDictSchema` is sufficient: + +```python +# Standalone schema - no class definition needed +schema = pa.TensorDictSchema( + keys={ + "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)), + "action": pa.Tensor(dtype=torch.int64, shape=(None,)), + }, + batch_size=(32,) +) + +validated = schema.validate(batch) +``` + +For reusable, self-documenting schemas with type hints, use `TensorDictModel`: + +```python +class MySchema(pa.TensorDictModel): + observation: torch.float32 = pa.Field(shape=(None, 10)) + action: torch.int64 = pa.Field(shape=(None,)) + + class Config: + batch_size = (32,) + +# Schema is automatically built from class definition +schema = MySchema.to_schema() # Same as standalone schema above! +validated = MySchema.validate(batch) # Direct validation +``` + +--- + + --- ## 5. Implementation Roadmap @@ -874,7 +989,9 @@ pandera/ │ ├── __init__.py │ ├── container.py # TensorDictSchema │ ├── components.py # Tensor -│ ├── model.py # TensorDictModel +│ ├── model.py # TensorDictModel +│ ├── model_components.py # Field, TensorDictFieldInfo +│ ├── model_config.py # BaseConfig │ └── types.py # Type aliases ├── backends/ │ └── tensordict/ @@ -884,8 +1001,5 @@ pandera/ │ └── register.py # Backend registration ├── engines/ │ └── tensordict_engine.py # Engine for torch dtype registry -├── typing/ -│ └── tensordict.py # Annotation types (TensorDict, TensorClass) └── tensordict.py # Entry point: import pandera.tensordict as pa - ``` diff --git a/tests/pandas/test_model_from_records_dicts.py b/tests/pandas/test_model_from_records_dicts.py index 98e4fd61d..aa89d9e7b 100644 --- a/tests/pandas/test_model_from_records_dicts.py +++ b/tests/pandas/test_model_from_records_dicts.py @@ -23,7 +23,7 @@ class Schema(DataFrameModel): {"state": "NY", "city": "New York", "price": 8.0}, {"state": "FL", "city": "Miami", "price": 12.0}, ] - + pandera_validated_df = DataFrame.from_records(Schema, raw_data) pandas_df = pd.DataFrame.from_records(raw_data) assert pandera_validated_df.equals(Schema.validate(pandas_df)) diff --git a/tests/pandas/test_pydantic.py b/tests/pandas/test_pydantic.py index e465625f3..a12803ffe 100644 --- a/tests/pandas/test_pydantic.py +++ b/tests/pandas/test_pydantic.py @@ -1,10 +1,6 @@ """Unit tests for pydantic compatibility.""" -from typing import ( - Generic, - Optional, - TypeVar, -) +from typing import Generic, Optional, TypeVar import pandas as pd import pytest diff --git a/tests/tensordict/__init__.py b/tests/tensordict/__init__.py new file mode 100644 index 000000000..7ff5dbd57 --- /dev/null +++ b/tests/tensordict/__init__.py @@ -0,0 +1 @@ +"""Unit tests for pandera tensordict integration.""" diff --git a/tests/tensordict/test_tensordict_container.py b/tests/tensordict/test_tensordict_container.py new file mode 100644 index 000000000..a0dc9f741 --- /dev/null +++ b/tests/tensordict/test_tensordict_container.py @@ -0,0 +1,667 @@ +"""Unit tests for TensorDict container and components.""" + +import pytest + +try: + import torch + from tensordict import TensorDict, tensorclass +except ImportError: + torch = None + TensorDict = None + tensorclass = None + +torch_condition = pytest.mark.skipif( + torch is None, reason="torch not installed" +) + + +@torch_condition +class TestTensorComponent: + """Tests for Tensor component.""" + + def test_tensor_creation(self): + """Test Tensor component creation.""" + from pandera.tensordict import Tensor + + tensor = Tensor(dtype=torch.float32, shape=(None, 10)) + assert tensor.dtype.type == torch.float32 + assert tensor.shape == (None, 10) + + def test_tensor_with_checks(self): + """Test Tensor component with checks.""" + from pandera import Check + from pandera.tensordict import Tensor + + tensor = Tensor( + dtype=torch.float32, + shape=(None, 10), + checks=Check.greater_than(0.0), + ) + assert tensor.dtype.type == torch.float32 + assert tensor.shape == (None, 10) + assert len(tensor.checks) == 1 + + def test_tensor_repr(self): + """Test Tensor repr.""" + from pandera.tensordict import Tensor + + tensor = Tensor(dtype=torch.float32, shape=(None, 10)) + repr_str = repr(tensor) + assert "Tensor" in repr_str + assert "float32" in repr_str + + +@torch_condition +class TestTensorDictSchema: + """Tests for TensorDictSchema.""" + + def test_tensordict_schema_creation(self): + """Test TensorDictSchema creation.""" + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(None, 10)), + "action": Tensor(dtype=torch.float32, shape=(None, 5)), + }, + batch_size=(32,), + ) + assert "observation" in schema.keys + assert "action" in schema.keys + assert schema.batch_size == (32,) + + def test_tensordict_schema_from_list(self): + """Test TensorDictSchema creation from key list.""" + from pandera.tensordict import TensorDictSchema + + schema = TensorDictSchema( + keys=["observation", "action"], batch_size=(32,) + ) + assert "observation" in schema.keys + assert "action" in schema.keys + + def test_tensordict_schema_repr(self): + """Test TensorDictSchema repr.""" + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32), + }, + batch_size=(32,), + ) + repr_str = repr(schema) + assert "TensorDictSchema" in repr_str + assert "observation" in repr_str + + +@torch_condition +class TestTensorDictValidation: + """Tests for TensorDict validation.""" + + def test_validate_valid_tensordict(self): + """Test validation of valid TensorDict.""" + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(32, 10)), + "action": Tensor(dtype=torch.float32, shape=(32, 5)), + }, + batch_size=(32,), + ) + + td = TensorDict( + { + "observation": torch.randn(32, 10), + "action": torch.randn(32, 5), + }, + batch_size=[32], + ) + + result = schema.validate(td) + assert isinstance(result, TensorDict) + + def test_validate_invalid_batch_size(self): + """Test validation fails with wrong batch size.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(32, 10)), + }, + batch_size=(32,), + ) + + td = TensorDict( + {"observation": torch.randn(16, 10)}, + batch_size=[16], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_validate_missing_key(self): + """Test validation fails with missing key.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32), + "action": Tensor(dtype=torch.float32), + }, + batch_size=(32,), + ) + + td = TensorDict( + {"observation": torch.randn(32, 10)}, + batch_size=[32], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_validate_wrong_dtype(self): + """Test validation fails with wrong dtype.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32), + }, + batch_size=(32,), + ) + + td = TensorDict( + {"observation": torch.randn(32, 10).to(torch.float64)}, + batch_size=[32], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_validate_wrong_shape(self): + """Test validation fails with wrong shape.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(32, 10)), + }, + batch_size=(32,), + ) + + td = TensorDict( + {"observation": torch.randn(32, 20)}, + batch_size=[32], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_validate_lazy(self): + """Test lazy validation collects all errors.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(None, 10)), + "action": Tensor(dtype=torch.float32, shape=(None, 5)), + }, + batch_size=(32,), + ) + + td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randn(16, 5)}, + batch_size=None, + ) + + with pytest.raises(errors.SchemaErrors) as exc_info: + schema.validate(td, lazy=True) + + assert len(exc_info.value.schema_errors) > 0 + + +@torch_condition +class TestTensorDictSchemaChecks: + """Tests for TensorDictSchema with value checks.""" + + def test_validate_with_value_check(self): + """Test validation with value check.""" + from pandera import Check + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "values": Tensor( + dtype=torch.float32, + shape=(None,), + checks=Check.greater_than(0.0), + ), + }, + batch_size=(10,), + ) + + td = TensorDict( + { + "values": torch.tensor( + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0] + ) + }, + batch_size=[10], + ) + + result = schema.validate(td) + assert isinstance(result, TensorDict) + + def test_validate_value_check_failure(self): + """Test validation fails with value check failure.""" + from pandera import Check, errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "values": Tensor( + dtype=torch.float32, + shape=(None,), + checks=Check.greater_than(0.0), + ), + }, + batch_size=(10,), + ) + + td = TensorDict( + { + "values": torch.tensor( + [1.0, -2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0] + ) + }, + batch_size=[10], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + +@torch_condition +class TestTensorDictSchemaBatchSize: + """Tests for batch size validation.""" + + def test_batch_size_with_none_dimension(self): + """Test batch size with None dimension allows any size.""" + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(None, 10)), + }, + batch_size=(None,), + ) + + td = TensorDict({"observation": torch.randn(5, 10)}, batch_size=[5]) + result = schema.validate(td) + assert isinstance(result, TensorDict) + + td = TensorDict( + {"observation": torch.randn(100, 10)}, batch_size=[100] + ) + result = schema.validate(td) + assert isinstance(result, TensorDict) + + def test_batch_size_exact_match(self): + """Test batch size must match exactly when specified.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(None, 10)), + }, + batch_size=(32,), + ) + + td = TensorDict({"observation": torch.randn(64, 10)}, batch_size=[64]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + +@torch_condition +class TestShapeValidation: + """Tests for shape validation.""" + + def test_shape_with_none_allows_any(self): + """Test shape with None allows any size for that dimension.""" + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "data": Tensor(dtype=torch.float32, shape=(None, None, 3)), + }, + batch_size=(10,), + ) + + td = TensorDict({"data": torch.randn(10, 32, 3)}, batch_size=[10]) + result = schema.validate(td) + assert isinstance(result, TensorDict) + + def test_shape_exact_match(self): + """Test shape must match exactly when specified.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "data": Tensor(dtype=torch.float32, shape=(32, 10)), + }, + batch_size=(32,), + ) + + td = TensorDict({"data": torch.randn(32, 20)}, batch_size=[32]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + +@torch_condition +class TestTensorClassValidation: + """Tests for tensorclass validation.""" + + def test_validate_tensorclass(self): + """Test validation of tensorclass object.""" + from pandera.tensordict import Tensor, TensorDictSchema + + @tensorclass + class TCData: + observation: torch.Tensor + action: torch.Tensor + + schema = TensorDictSchema( + keys={ + "observation": Tensor(dtype=torch.float32, shape=(32, 10)), + "action": Tensor(dtype=torch.float32, shape=(32, 5)), + }, + batch_size=(32,), + ) + + tc = TCData( + observation=torch.randn(32, 10), + action=torch.randn(32, 5), + batch_size=[32], + ) + + result = schema.validate(tc) + assert result is not None + + +@torch_condition +class TestTensorDictErrorCases: + """Comprehensive error case tests.""" + + def test_invalid_input_type(self): + """Test validation fails with non-TensorDict input.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema(keys={"x": Tensor(dtype=torch.float32)}) + + with pytest.raises( + errors.BackendNotFoundError, match="Backend not found" + ): + schema.validate({"x": torch.randn(10)}) + + def test_invalid_tensor_in_tensordict(self): + """Test validation fails when key contains non-tensor.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10,), + ) + + td = TensorDict({"x": "not a tensor"}, batch_size=[10]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_batch_size_length_mismatch(self): + """Test validation fails when batch_size dimensions don't match.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10, 5), + ) + + td = TensorDict({"x": torch.randn(10)}, batch_size=[10]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_batch_size_value_mismatch_first_dim(self): + """Test validation fails when first batch_size dimension doesn't match.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10,), + ) + + td = TensorDict({"x": torch.randn(5)}, batch_size=[5]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_batch_size_value_mismatch_second_dim(self): + """Test validation fails when second batch_size dimension doesn't match.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10, 5), + ) + + td = TensorDict({"x": torch.randn(10, 3)}, batch_size=[10, 3]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_multiple_missing_keys(self): + """Test validation fails with multiple missing keys.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "a": Tensor(dtype=torch.float32), + "b": Tensor(dtype=torch.float32), + "c": Tensor(dtype=torch.float32), + }, + batch_size=(10,), + ) + + td = TensorDict({"a": torch.randn(10)}, batch_size=[10]) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_multiple_dtype_mismatches(self): + """Test validation fails with multiple dtype errors.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "a": Tensor(dtype=torch.float32), + "b": Tensor(dtype=torch.float32), + }, + batch_size=(10,), + ) + + td = TensorDict( + { + "a": torch.randn(10).to(torch.int32), + "b": torch.randn(10).to(torch.int64), + }, + batch_size=[10], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_multiple_shape_errors(self): + """Test validation fails with multiple shape errors.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "a": Tensor(dtype=torch.float32, shape=(10, 5)), + "b": Tensor(dtype=torch.float32, shape=(10, 3)), + }, + batch_size=(10,), + ) + + td = TensorDict( + {"a": torch.randn(10, 7), "b": torch.randn(10, 8)}, + batch_size=[10], + ) + + with pytest.raises(errors.SchemaError): + schema.validate(td) + + def test_lazy_collects_all_errors(self): + """Test lazy validation collects multiple errors.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "a": Tensor(dtype=torch.float32, shape=(32, 5)), + "b": Tensor(dtype=torch.float32, shape=(32, 3)), + "c": Tensor(dtype=torch.float32), + }, + batch_size=(32,), + ) + + td = TensorDict( + { + "a": torch.randn(16, 5), + "b": torch.randn(16, 8), + }, + batch_size=[16], + ) + + with pytest.raises(errors.SchemaErrors) as exc_info: + schema.validate(td, lazy=True) + + assert len(exc_info.value.schema_errors) >= 2 + + def test_error_message_content_batch_size(self): + """Test error message contains batch_size info.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10,), + ) + + td = TensorDict({"x": torch.randn(5)}, batch_size=[5]) + + try: + schema.validate(td) + except errors.SchemaError as e: + error_messages = [str(e)] + assert any("batch_size" in msg for msg in error_messages) + + def test_error_message_content_dtype(self): + """Test error message contains dtype info.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10,), + ) + + td = TensorDict( + {"x": torch.randn(10).to(torch.int32)}, batch_size=[10] + ) + + try: + schema.validate(td) + except errors.SchemaError as e: + error_messages = [str(e)] + assert any("dtype" in msg for msg in error_messages) + + def test_error_message_content_shape(self): + """Test error message contains shape info.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32, shape=(10, 5))}, + batch_size=(10,), + ) + + td = TensorDict({"x": torch.randn(10, 3)}, batch_size=[10]) + + try: + schema.validate(td) + except errors.SchemaError as e: + error_messages = [str(e)] + assert any("shape" in msg for msg in error_messages) + + def test_error_message_content_missing_key(self): + """Test error message contains missing key info.""" + from pandera import errors + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={ + "x": Tensor(dtype=torch.float32), + "y": Tensor(dtype=torch.float32), + }, + batch_size=(10,), + ) + + td = TensorDict({"x": torch.randn(10)}, batch_size=[10]) + + try: + schema.validate(td) + except errors.SchemaError as e: + error_messages = [str(e)] + assert any("Missing key" in msg for msg in error_messages) + + def test_schema_error_reason_codes(self): + """Test that error reason codes are correctly set.""" + from pandera import errors + from pandera.errors import SchemaErrorReason + from pandera.tensordict import Tensor, TensorDictSchema + + schema = TensorDictSchema( + keys={"x": Tensor(dtype=torch.float32)}, + batch_size=(10,), + ) + + td = TensorDict({"x": torch.randn(5)}, batch_size=[5]) + + try: + schema.validate(td) + except errors.SchemaError as e: + reason_codes = [e.reason_code] + assert SchemaErrorReason.WRONG_DATATYPE in reason_codes + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/tensordict/test_tensordict_engine.py b/tests/tensordict/test_tensordict_engine.py new file mode 100644 index 000000000..5fd4965d1 --- /dev/null +++ b/tests/tensordict/test_tensordict_engine.py @@ -0,0 +1,149 @@ +"""Unit tests for TensorDict engine.""" + +import pytest + +try: + import torch +except ImportError: + torch = None + +torch_condition = pytest.mark.skipif( + torch is None, reason="torch not installed" +) + + +@torch_condition +class TestTensorDictEngine: + """Tests for TensorDict engine.""" + + def test_engine_import(self): + """Test that tensordict_engine can be imported.""" + from pandera.engines import tensordict_engine + + assert tensordict_engine is not None + + def test_engine_dtype_from_string(self): + """Test engine dtype from string.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.Engine.dtype("float32") + assert dtype is not None + + def test_engine_dtype_from_torch_dtype(self): + """Test engine dtype from torch.dtype.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.Engine.dtype(torch.float32) + assert dtype is not None + + def test_engine_dtype_invalid(self): + """Test engine raises on invalid dtype.""" + from pandera.engines import tensordict_engine + + with pytest.raises(ValueError): + tensordict_engine.Engine.dtype("invalid_dtype") + + def test_datatype_coerce(self): + """Test DataType coerce method.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.DataType(torch.float32) + tensor = torch.randn(10).to(torch.float64) + coerced = dtype.coerce(tensor) + assert coerced.dtype == torch.float32 + + def test_datatype_coerce_value(self): + """Test DataType coerce_value method.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.DataType(torch.float32) + value = dtype.coerce_value(1.0) + assert isinstance(value, torch.Tensor) + + def test_datatype_str(self): + """Test DataType str representation.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.DataType(torch.float32) + assert "float32" in str(dtype) + + def test_datatype_repr(self): + """Test DataType repr.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.DataType(torch.float32) + assert "DataType" in repr(dtype) + + def test_engine_multiple_dtypes(self): + """Test engine with multiple dtypes.""" + from pandera.engines import tensordict_engine + + dtypes = [ + "float32", + "float64", + "int32", + "int64", + "int16", + "int8", + "uint8", + "bool", + ] + + for dtype_str in dtypes: + dtype = tensordict_engine.Engine.dtype(dtype_str) + assert dtype is not None + + +@torch_condition +class TestTensorDictEngineErrorCases: + """Error case tests for tensordict engine.""" + + def test_engine_dtype_from_invalid_string(self): + """Test engine raises on invalid dtype string.""" + from pandera.engines import tensordict_engine + + with pytest.raises(ValueError, match="not understood"): + tensordict_engine.Engine.dtype("invalid_type") + + def test_engine_dtype_from_invalid_type(self): + """Test engine raises on invalid dtype type.""" + from pandera.engines import tensordict_engine + + with pytest.raises(ValueError, match="not understood"): + tensordict_engine.Engine.dtype(123) + + def test_datatype_coerce_wrong_type(self): + """Test DataType coerce raises on wrong type.""" + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.DataType(torch.float32) + + with pytest.raises(Exception): + dtype.coerce("not a tensor") + + def test_datatype_try_coerce_wrong_type(self): + """Test DataType try_coerce raises on wrong type.""" + from pandera import errors + from pandera.engines import tensordict_engine + + dtype = tensordict_engine.DataType(torch.float32) + + with pytest.raises(errors.ParserError): + dtype.try_coerce("not a tensor") + + +@torch_condition +class TestTensorDictEngineNotInstalled: + """Tests when torch is not installed.""" + + def test_no_torch_placeholder(self): + """Test that engine has placeholder when torch not installed.""" + import pandera.engines.tensordict_engine as engine + + if torch is None: + assert engine.DataType is None + assert engine.Engine is None + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py new file mode 100644 index 000000000..c0c5da916 --- /dev/null +++ b/tests/tensordict/test_tensordict_model.py @@ -0,0 +1,145 @@ +"""Unit tests for TensorDictModel.""" + +import pytest + +try: + import torch + from tensordict import TensorDict + + from pandera.tensordict import DataType # type: ignore[attr-defined] + + _DataType = DataType # type: ignore[misc] +except ImportError: + torch = None + TensorDict = None + DataType = None # type: ignore[misc, assignment] + +torch_condition = pytest.mark.skipif( + torch is None, reason="torch not installed" +) + + +@torch_condition +class TestTensorDictModelBasic: + """Tests for basic TensorDictModel functionality.""" + + def test_model_with_field_annotations(self): + """Test model with field annotations.""" + from pandera.tensordict import Field, TensorDictModel + + class MyModel(TensorDictModel): + observation: torch.float32 = Field(shape=(None, 10)) + action: torch.int64 = Field(shape=(None,)) + + class Config: + batch_size = (32,) + + schema = MyModel.to_schema() + assert "observation" in schema.keys + assert "action" in schema.keys + + def test_model_validation(self): + """Test model validation with valid data.""" + from pandera.tensordict import Field, TensorDictModel + + class MyModel(TensorDictModel): + observation: torch.float32 = Field(shape=(None, 10)) + + class Config: + batch_size = (32,) + + batch = TensorDict( + {"observation": torch.randn(32, 10)}, + batch_size=[32], + ) + + validated = MyModel.validate(batch) + assert isinstance(validated, TensorDict) + + def test_model_validation_failure(self): + """Test model validation fails with invalid data.""" + import pandera.errors as errors + from pandera.tensordict import Field, TensorDictModel + + class MyModel(TensorDictModel): + observation: torch.float32 = Field(shape=(None, 10)) + + class Config: + batch_size = (32,) + + batch = TensorDict( + {"observation": torch.randn(16, 10)}, + batch_size=[16], + ) + + with pytest.raises(errors.SchemaError): + MyModel.validate(batch) + + +@torch_condition +class TestTensorDictModelWithChecks: + """Tests for TensorDictModel with value checks.""" + + def test_model_with_field_checks(self): + """Test model with Field-level checks.""" + import pandera.errors as errors + from pandera import Check + from pandera.tensordict import Field, TensorDictModel + + class MyModel(TensorDictModel): + values: torch.float32 = Field( + shape=(None,), + ge=0.0, + le=1.0, + ) + + class Config: + batch_size = (10,) + + # Valid data + batch = TensorDict( + {"values": torch.rand(10)}, + batch_size=[10], + ) + validated = MyModel.validate(batch) + assert isinstance(validated, TensorDict) + + # Invalid data - outside range + batch_invalid = TensorDict( + {"values": torch.randn(10)}, # Some negative values + batch_size=[10], + ) + + with pytest.raises(errors.SchemaError): + MyModel.validate(batch_invalid) + + +@torch_condition +class TestTensorDictModelInheritance: + """Tests for TensorDictModel inheritance.""" + + def test_model_inheritance(self): + """Test model can be inherited from.""" + from pandera.tensordict import Field, TensorDictModel + + class BaseModel(TensorDictModel): + observation: torch.float32 = Field(shape=(None, 10)) + + class Config: + batch_size = (32,) + + class ChildModel(BaseModel): + action: torch.int64 = Field(shape=(None,)) + + # Base schema should still work + base_schema = BaseModel.to_schema() + assert "observation" in base_schema.keys + + # Child schema should have both fields + child_schema = ChildModel.to_schema() + assert "observation" in child_schema.keys + assert "action" in child_schema.keys + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/xarray/test_cf_xarray.py b/tests/xarray/test_cf_xarray.py index c31091eb9..0f4222f48 100644 --- a/tests/xarray/test_cf_xarray.py +++ b/tests/xarray/test_cf_xarray.py @@ -6,12 +6,7 @@ xr = pytest.importorskip("xarray") import pandera.errors -from pandera.xarray import ( - Check, - DataArraySchema, - DatasetSchema, - DataVar, -) +from pandera.xarray import Check, DataArraySchema, DatasetSchema, DataVar class TestCfStandardName: diff --git a/tests/xarray/test_data_array_schema.py b/tests/xarray/test_data_array_schema.py index ef62314f8..834825764 100644 --- a/tests/xarray/test_data_array_schema.py +++ b/tests/xarray/test_data_array_schema.py @@ -326,9 +326,7 @@ class TestDataArrayCheckMethods: @pytest.fixture() def backend(self): - from pandera.backends.xarray.container import ( - DataArraySchemaBackend, - ) + from pandera.backends.xarray.container import DataArraySchemaBackend return DataArraySchemaBackend() @@ -515,10 +513,7 @@ def test_check_strict_coords_fail(self, backend): assert any(not r.passed for r in results) def test_schema_scope_checks_skipped_with_data_only(self, backend): - from pandera.config import ( - ValidationDepth, - config_context, - ) + from pandera.config import ValidationDepth, config_context da = xr.DataArray(np.zeros(2), dims="x", name="wrong") schema = DataArraySchema(name="correct") diff --git a/tests/xarray/test_dataset_schema.py b/tests/xarray/test_dataset_schema.py index 685055d0c..7b829cab8 100644 --- a/tests/xarray/test_dataset_schema.py +++ b/tests/xarray/test_dataset_schema.py @@ -6,12 +6,7 @@ xr = pytest.importorskip("xarray") import pandera.errors -from pandera.xarray import ( - Check, - Coordinate, - DatasetSchema, - DataVar, -) +from pandera.xarray import Check, Coordinate, DatasetSchema, DataVar def test_data_vars_and_coords(): @@ -194,9 +189,7 @@ class TestDatasetCheckMethods: @pytest.fixture() def backend(self): - from pandera.backends.xarray.container import ( - DatasetSchemaBackend, - ) + from pandera.backends.xarray.container import DatasetSchemaBackend return DatasetSchemaBackend() @@ -460,10 +453,7 @@ def test_check_broadcastable_with_fail(self, backend): assert not results or all(r.passed for r in results) def test_schema_scope_checks_skipped_with_data_only(self, backend): - from pandera.config import ( - ValidationDepth, - config_context, - ) + from pandera.config import ValidationDepth, config_context ds = xr.Dataset({"a": (["x"], np.zeros(2))}) schema = DatasetSchema(data_vars={"a": DataVar()}, dims=("y",)) diff --git a/tests/xarray/test_encoding.py b/tests/xarray/test_encoding.py index 59b68e6e4..b6b810d8a 100644 --- a/tests/xarray/test_encoding.py +++ b/tests/xarray/test_encoding.py @@ -6,12 +6,7 @@ xr = pytest.importorskip("xarray") import pandera.errors -from pandera.xarray import ( - Check, - DataArraySchema, - DatasetSchema, - DataVar, -) +from pandera.xarray import Check, DataArraySchema, DatasetSchema, DataVar # ------------------------------------------------------------------ # Helpers