From 6e582beb8d0bfe46794432ddfa24627e693b0a87 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 20:57:48 -0400 Subject: [PATCH 01/27] feat(tensordict): implement Phase 1 core infrastructure and API - Add TensorDictSchema and Tensor component classes - Implement TensorDictModel with class-based schema definition - Add Field descriptor with check parameters (gt, ge, lt, le, etc.) - Create tensordict_engine for PyTorch dtype resolution - Implement TensorDictSchemaBackend with batch_size, key, dtype, shape validation - Add TensorDictCheckBackend for tensor value checks - Update spec examples to use correct Field API pattern - Add module layout documentation to spec --- pandera/api/tensordict/__init__.py | 8 + pandera/api/tensordict/components.py | 66 ++++ pandera/api/tensordict/container.py | 139 +++++++++ pandera/api/tensordict/model.py | 176 +++++++++++ pandera/api/tensordict/model_components.py | 160 ++++++++++ pandera/api/tensordict/model_config.py | 14 + pandera/api/tensordict/types.py | 26 ++ pandera/backends/tensordict/__init__.py | 1 + pandera/backends/tensordict/base.py | 340 +++++++++++++++++++++ pandera/backends/tensordict/checks.py | 78 +++++ pandera/backends/tensordict/register.py | 47 +++ pandera/engines/tensordict_engine.py | 124 ++++++++ pandera/tensordict_api.py | 8 + specs/pytorch-tensordict.md | 27 +- 14 files changed, 1203 insertions(+), 11 deletions(-) create mode 100644 pandera/api/tensordict/__init__.py create mode 100644 pandera/api/tensordict/components.py create mode 100644 pandera/api/tensordict/container.py create mode 100644 pandera/api/tensordict/model.py create mode 100644 pandera/api/tensordict/model_components.py create mode 100644 pandera/api/tensordict/model_config.py create mode 100644 pandera/api/tensordict/types.py create mode 100644 pandera/backends/tensordict/__init__.py create mode 100644 pandera/backends/tensordict/base.py create mode 100644 pandera/backends/tensordict/checks.py create mode 100644 pandera/backends/tensordict/register.py create mode 100644 pandera/engines/tensordict_engine.py create mode 100644 pandera/tensordict_api.py diff --git a/pandera/api/tensordict/__init__.py b/pandera/api/tensordict/__init__.py new file mode 100644 index 000000000..ec48b2eb7 --- /dev/null +++ b/pandera/api/tensordict/__init__.py @@ -0,0 +1,8 @@ +"""Pandera TensorDict API.""" + +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"] diff --git a/pandera/api/tensordict/components.py b/pandera/api/tensordict/components.py new file mode 100644 index 000000000..48b518830 --- /dev/null +++ b/pandera/api/tensordict/components.py @@ -0,0 +1,66 @@ +"""Tensor component for TensorDict validation.""" + +from __future__ import annotations + +from typing import Any + +from pandera.api.base.types import CheckList +from pandera.api.dataframe.components import ComponentSchema + +try: + import torch +except ImportError: + torch = None + + +class Tensor(ComponentSchema): + """Schema component for validating a tensor in a TensorDict.""" + + 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, + title: str | None = None, + description: str | None = None, + metadata: dict | None = None, + drop_invalid_rows: bool = False, + ) -> None: + """Create a tensor schema component. + + :param dtype: Expected torch dtype for the tensor. + :param shape: Expected shape for the tensor. Use None for dimensions + that can be any size. The first dimension should typically be None + to allow for arbitrary batch sizes. + :param checks: Checks to verify validity of the tensor values. + :param nullable: Whether or not the tensor can be None. + :param coerce: If True, coerce the tensor to the specified dtype. + :param name: Name of the tensor in the TensorDict. + :param title: A human-readable label for the tensor. + :param description: An arbitrary textual description of the tensor. + :param metadata: An optional key-value data. + :param drop_invalid_rows: Not applicable for TensorDict. + + :raises SchemaInitError: if impossible to build schema from parameters + """ + super().__init__( + dtype=dtype, + checks=checks, + nullable=nullable, + coerce=coerce, + name=name, + title=title, + description=description, + metadata=metadata, + drop_invalid_rows=drop_invalid_rows, + ) + self.shape = shape + + def __repr__(self) -> str: + return ( + f"Tensor(dtype={self.dtype}, shape={self.shape}, " + f"nullable={self.nullable})" + ) diff --git a/pandera/api/tensordict/container.py b/pandera/api/tensordict/container.py new file mode 100644 index 000000000..e477136a1 --- /dev/null +++ b/pandera/api/tensordict/container.py @@ -0,0 +1,139 @@ +"""TensorDict Schema for pandera.""" + +from __future__ import annotations + +import sys +from typing import Any + +from pandera.api.base.types import CheckList +from pandera.api.dataframe.container import DataFrameSchema as _DataFrameSchema +from pandera.api.tensordict.components import Tensor +from pandera.api.tensordict.types import TensorDictCheckObjects +from pandera.config import config_context, get_config_context + +if sys.version_info < (3, 11): + from typing_extensions import Self +else: + from typing import Self + +try: + import torch +except ImportError: + torch = None + + +class TensorDictSchema(_DataFrameSchema[TensorDictCheckObjects]): + """A schema for validating TensorDict and tensorclass objects.""" + + def __init__( + self, + keys: dict[str, Tensor] | list[str] | None = None, + batch_size: tuple[int, ...] | None = None, + dtype: Any = None, + checks: CheckList | None = None, + coerce: bool = False, + nullable: bool = False, + name: str | None = None, + title: str | None = None, + description: str | None = None, + metadata: dict | None = None, + drop_invalid_rows: bool = False, + ) -> None: + """Create a TensorDict schema. + + :param keys: A dictionary mapping key names to Tensor schema objects, + or a list of key names (all tensors will be validated for dtype only). + :param batch_size: The expected batch size. Use None in dimensions + to allow for any size. + :param dtype: Expected dtype for all tensors in the container. + :param checks: Checks to apply to the entire TensorDict. + :param coerce: If True, coerce tensors to the specified dtype. + :param nullable: Whether or not tensors can be None. + :param name: Name of the schema. + :param title: A human-readable label for the schema. + :param description: An arbitrary textual description of the schema. + :param metadata: An optional key-value data. + :param drop_invalid_rows: Not applicable for TensorDict. + + :raises SchemaInitError: if impossible to build schema from parameters + """ + if keys is None: + keys = {} + if isinstance(keys, list): + keys = {k: Tensor(dtype=dtype) for k in keys} + + super().__init__( + columns=keys, + checks=checks, + coerce=coerce, + name=name, + title=title, + description=description, + metadata=metadata, + drop_invalid_rows=drop_invalid_rows, + ) + self.batch_size = batch_size + self._dtype = dtype + + @staticmethod + def register_default_backends(check_obj_cls: type): + from pandera.backends.tensordict.register import ( + register_tensordict_backends, + ) + + register_tensordict_backends() + + @property + def dtype(self): + return self._dtype + + @dtype.setter + def dtype(self, value) -> None: + if torch is None: + self._dtype = value + else: + from pandera.engines import tensordict_engine + self._dtype = tensordict_engine.Engine.dtype(value) if value else None + + def validate( + self, + check_obj: TensorDictCheckObjects, + lazy: bool = False, + inplace: bool = False, + ) -> TensorDictCheckObjects: + """Validate a TensorDict or tensorclass against the schema. + + :param check_obj: TensorDict or tensorclass to validate. + :param lazy: If True, collect all errors. Otherwise, fail fast. + :param inplace: If True, modify the input in place. + :returns: Validated TensorDict or tensorclass. + + :raises SchemaError: when TensorDict violates schema. + """ + if not get_config_context().validation_enabled: + return check_obj + + with config_context(validation_depth=None): + output = self.get_backend(check_obj).validate( + check_obj=check_obj, + schema=self, + lazy=lazy, + inplace=inplace, + ) + + return output + + @_DataFrameSchema.dtype.setter # type: ignore[attr-defined] + def dtype(self, value) -> None: + """Set the dtype property.""" + if torch is None: + self._dtype = value + else: + from pandera.engines import tensordict_engine + self._dtype = tensordict_engine.Engine.dtype(value) if value else None + + def __repr__(self) -> str: + return ( + f"TensorDictSchema(keys={list(self.columns.keys())}, " + f"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..470bb959c --- /dev/null +++ b/pandera/api/tensordict/model.py @@ -0,0 +1,176 @@ +"""Declarative TensorDict models.""" + +import sys +from typing import Any, ClassVar, Union, cast + +from pandera.api.base.model import BaseModel +from pandera.api.checks import Check +from pandera.api.dataframe.model_components import ( + CHECK_KEY, + Field, + FieldInfo, +) +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 + +_CONFIG_KEY = "Config" + +try: + from typing_extensions import Self +except ImportError: + from typing import Self + + +def _is_field(name: str) -> bool: + return not name.startswith("_") and name != _CONFIG_KEY + + +class TensorDictModel(BaseModel): + """Model of a TensorDict schema. + + *new in 0.19.0* + + See the :ref:`User Guide ` for more. + """ + + Config: type[BaseConfig] = BaseConfig + + __fields__: ClassVar[dict[str, tuple[AnnotationInfo, FieldInfo]]] = {} + + @classmethod + def build_schema_(cls, **kwargs) -> TensorDictSchema: + columns = cls._build_columns( + cls.__fields__, + cls.__checks__, + ) + return TensorDictSchema( + keys=columns, + checks=cls.__root_checks__, + batch_size=cls.__config__.batch_size, + name=cls.__config__.name, + title=cls.__config__.title, + description=cls.__config__.description, + coerce=cls.__config__.coerce, + **kwargs, + ) + + @classmethod + def _build_columns( + cls, + fields: dict[str, tuple[AnnotationInfo, FieldInfo]], + checks: dict[str, list[Check]], + ) -> dict[str, Tensor]: + columns = {} + for name, (annotation, field) in fields.items(): + dtype = annotation.arg + if dtype is None: + raise SchemaInitError( + f"expected annotation for field '{name}'" + ) + + field_checks = checks.get(name, []) + column_kwargs = field.to_tensor_kwargs( + dtype, + optional=not annotation.optional, + checks=field_checks, + ) + column = Tensor(**column_kwargs) + columns[name] = column + + return columns + + @classmethod + def _collect_fields(cls) -> dict[str, tuple[AnnotationInfo, FieldInfo]]: + import inspect + + annotations = inspect.get_annotations(cls) + attrs = vars(cls) + + 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.get(field_name) + if field is None: + field = TensorDictField() + elif CHECK_KEY in vars(field): + field = field[CHECK_KEY] + fields[field_name] = (AnnotationInfo(annotation), field) + return fields + + @classmethod + def _collect_config_and_extras(cls): + """Collect config options from class.""" + if "Config" in cls.__dict__: + cls.Config.name = ( + cls.Config.name + if hasattr(cls.Config, "name") + else cls.__name__ + ) + else: + cls.Config = type("Config", (cls.Config,), {"name": cls.__name__}) + + config_options = {} + extras = {} + for name, value in vars(cls.Config).items(): + if _is_field(name): + config_options[name] = value + elif _is_field(name): + extras[name] = value + + return cls.Config, extras + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + import inspect + + subclass_annotations = inspect.get_annotations(cls) + + for field_name in subclass_annotations.keys(): + if _is_field(field_name) and field_name not in cls.__dict__: + field = TensorDictField() + field.__set_name__(cls, field_name) + setattr(cls, field_name, field) + + cls.__config__, cls.__extras__ = cls._collect_config_and_extras() + cls.__fields__ = cls._collect_fields() + cls.__checks__ = {} + cls.__root_checks__ = [] + + @classmethod + def to_schema(cls, **kwargs) -> TensorDictSchema: + return cls.build_schema_(**kwargs) + + def validate(self, check_obj, *args, **kwargs): + schema = self.to_schema() + return schema.validate(check_obj, *args, **kwargs) + + def __new__(cls, *args, **kwargs) -> Any: + if cls is TensorDictModel: + raise NotImplementedError( + "TensorDictModel cannot be instantiated directly. " + "Create a subclass with field definitions." + ) + return super().__new__(cls) + + +def TensorDictField(): + """Field specification for TensorDict models.""" + from pandera.api.tensordict.model_components import Field + + return Field() diff --git a/pandera/api/tensordict/model_components.py b/pandera/api/tensordict/model_components.py new file mode 100644 index 000000000..68f28d17e --- /dev/null +++ b/pandera/api/tensordict/model_components.py @@ -0,0 +1,160 @@ +"""Field descriptors for TensorDict declarative models.""" + +from collections.abc import Iterable +from typing import Any, Union + +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__" + + +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`.""" + req = self.required and not optional + return self._get_schema_properties( + dtype, + checks=checks, + required=req, + shape=self.shape, + alias=self.alias, + title=self.title, + description=self.description, + default=self.default, + metadata=self.metadata, + nullable=self.nullable, + coerce=self.coerce, + ) + + +def Field( + *, + 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, + shape: tuple[int | None, ...] | None = None, + **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..1421c7457 --- /dev/null +++ b/pandera/api/tensordict/types.py @@ -0,0 +1,26 @@ +"""Type definitions for pandera tensordict integration.""" + +from typing import TYPE_CING, Any, Union + +import torch + +from pandera.api.checks import Check + +if TYPE_CING: + from tensordict import TensorDict, tensorclass + +TensorDictCheckObjects = Any + +TensorDictInputType = Union["TensorDict", "tensorclass", torch.Tensor] + +TensorDtypeInputTypes = Union[torch.dtype, str, type] + + +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/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..b8f1f1f5a --- /dev/null +++ b/pandera/backends/tensordict/base.py @@ -0,0 +1,340 @@ +"""TensorDict parsing, validation, and error-reporting backends.""" + +import copy +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, tensorclass + + return isinstance(obj, (TensorDict, tensorclass)) + except ImportError: + return False + + +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, + *, + lazy: bool = False, + inplace: bool = False, + ): + """Validate a TensorDict against the schema.""" + error_handler = ErrorHandler(lazy) + + if not _is_tensordict(check_obj): + raise TypeError( + f"Expected TensorDict or tensorclass, got {type(check_obj)}" + ) + + 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, + schema_context={"batch_size": actual_batch_size}, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_TYPE, + ) + 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, + schema_context={"batch_size": actual_batch_size}, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_TYPE, + ) + 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.columns: + if key not in check_obj: + error_msg = f"Missing key '{key}' in TensorDict" + error = SchemaError( + schema=schema, + schema_context={"keys": list(check_obj.keys())}, + message=error_msg, + reason_code=SchemaErrorReason.MISSING_COLUMN, + ) + 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.columns.items(): + if key not in check_obj: + continue + + tensor = check_obj[key] + + if not isinstance(tensor, torch.Tensor): + error_msg = f"Key '{key}' is not a torch.Tensor" + error = SchemaError( + schema=schema, + schema_context={"key": key, "type": type(tensor)}, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_TYPE, + ) + 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 + if tensor.dtype != expected_dtype: + error_msg = ( + f"Key '{key}': expected dtype {expected_dtype}, " + f"got {tensor.dtype}" + ) + error = SchemaError( + schema=schema, + schema_context={ + "key": key, + "dtype": str(tensor.dtype), + }, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_TYPE, + ) + 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, + schema_context={ + "key": key, + "shape": actual_shape, + }, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_TYPE, + ) + 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, + schema_context={ + "key": key, + "shape": actual_shape, + }, + message=error_msg, + reason_code=SchemaErrorReason.WRONG_TYPE, + ) + 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.""" + for key, tensor_schema in schema.columns.items(): + if key not in check_obj: + continue + + if not tensor_schema.checks: + continue + + tensor = check_obj[key] + is_tensorclass = hasattr(tensor, "as_dict") + tensor_data = ( + tensor.as_dict() if is_tensorclass else {key: tensor} + ) + + for check_index, check in enumerate(tensor_schema.checks): + try: + check_result = check(tensor_data, key) + except Exception as exc: + check_result = CoreCheckResult( + passed=False, + check=check, + check_index=check_index, + reason_code=SchemaErrorReason.CHECK_FAILED, + message=str(exc), + ) + + if not check_result.passed: + error_msg = ( + f"Check '{check}' failed for key '{key}': " + f"{check_result.message}" + ) + error = SchemaError( + schema=schema, + schema_context={"key": key}, + message=error_msg, + check=check, + check_index=check_index, + reason_code=SchemaErrorReason.CHECK_FAILED, + ) + 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.""" + return None + + 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/checks.py b/pandera/backends/tensordict/checks.py new file mode 100644 index 000000000..4cfbf970b --- /dev/null +++ b/pandera/backends/tensordict/checks.py @@ -0,0 +1,78 @@ +"""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 + +try: + import torch +except ImportError: + torch = None + + +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 apply(self, check_obj: Any): + """Apply the check function to the tensor data.""" + if self.check_fn is None: + return CheckResult(check_passed=True, failure_cases=None) + + if self.check.element_wise: + return self._apply_element_wise(check_obj) + else: + return self._apply_full_tensor(check_obj) + + def _apply_element_wise(self, check_obj: Any) -> CheckResult: + """Apply check element-wise to tensor.""" + if isinstance(check_obj, dict): + for key, tensor in check_obj.items(): + result = self.check_fn(tensor) + if isinstance(result, torch.Tensor): + passed = result.all().item() + else: + passed = bool(result) + if not passed: + return CheckResult( + check_passed=False, + failure_cases={"key": key, "failed_indices": []}, + ) + return CheckResult(check_passed=True, failure_cases=None) + return CheckResult(check_passed=True, failure_cases=None) + + def _apply_full_tensor(self, check_obj: Any) -> CheckResult: + """Apply check to full tensor.""" + if isinstance(check_obj, dict): + for key, tensor in check_obj.items(): + result = self.check_fn(tensor) + if isinstance(result, torch.Tensor): + passed = result.all().item() + else: + passed = bool(result) + if not passed: + return CheckResult( + check_passed=False, + failure_cases={"key": key}, + ) + return CheckResult(check_passed=True, failure_cases=None) + return CheckResult(check_passed=True, failure_cases=None) + + def postprocess(self, check_obj, check_output): + """Postprocesses the result of applying the check function.""" + if isinstance(check_output, CheckResult): + return check_output + return CheckResult(check_passed=bool(check_output), failure_cases=None) diff --git a/pandera/backends/tensordict/register.py b/pandera/backends/tensordict/register.py new file mode 100644 index 000000000..415263855 --- /dev/null +++ b/pandera/backends/tensordict/register.py @@ -0,0 +1,47 @@ +"""Register TensorDict backends.""" + +from functools import lru_cache +from typing import TYPE_CING, Any + +if TYPE_CING: + from tensordict import TensorDict, tensorclass + + +@lru_cache +def register_tensordict_backends( + check_cls_fqn: str | None = None, +): + """Register TensorDict backends. + + This function is called at schema initialization in the _register_*_backends + method. + """ + try: + from tensordict import TensorDict, tensorclass + except ImportError: + return + + from pandera.api.tensordict.container import TensorDictSchema + from pandera.api.tensordict.components import Tensor + from pandera.api.checks import Check + from pandera.backends.tensordict.base import TensorDictSchemaBackend + from pandera.backends.tensordict.checks import TensorDictCheckBackend + + TensorDictSchema.register_backend( + TensorDict, TensorDictSchemaBackend + ) + TensorDictSchema.register_backend( + tensorclass, TensorDictSchemaBackend + ) + Tensor.register_backend( + TensorDict, TensorDictSchemaBackend + ) + Tensor.register_backend( + tensorclass, TensorDictSchemaBackend + ) + Check.register_backend( + TensorDict, TensorDictCheckBackend + ) + Check.register_backend( + tensorclass, TensorDictCheckBackend + ) diff --git a/pandera/engines/tensordict_engine.py b/pandera/engines/tensordict_engine.py new file mode 100644 index 000000000..f968444b2 --- /dev/null +++ b/pandera/engines/tensordict_engine.py @@ -0,0 +1,124 @@ +"""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__() + object.__setattr__(self, "type", torch.dtype(dtype)) + + def __post_init__(self): + object.__setattr__(self, "type", torch.dtype(self.type)) + + 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 self.type(value) + + 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.""" + try: + return engine.Engine.dtype(cls, data_type) + except TypeError: + try: + torch_dtype = torch.dtype(data_type) + return DataType(torch_dtype) + except TypeError: + 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 = {} + for source_dtype in (pandera_dtype_cls.type,): + if source_dtype is not None: + equivalents[source_dtype] = pandera_dtype_cls + + 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(torch_dtype)) + except (AttributeError, ValueError): + pass + + _register_torch_dtypes() + +else: + DataType = None + Engine = None diff --git a/pandera/tensordict_api.py b/pandera/tensordict_api.py new file mode 100644 index 000000000..ec48b2eb7 --- /dev/null +++ b/pandera/tensordict_api.py @@ -0,0 +1,8 @@ +"""Pandera TensorDict API.""" + +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"] diff --git a/specs/pytorch-tensordict.md b/specs/pytorch-tensordict.md index 8f8001ebb..f3708b75e 100644 --- a/specs/pytorch-tensordict.md +++ b/specs/pytorch-tensordict.md @@ -163,20 +163,21 @@ For a more Pydantic-like experience, use class-based models: ```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( + logits: torch.Tensor = pa.Field( dtype=torch.float32, shape=(None, 10), - checks=Check.in_range(-5.0, 5.0) + ge=-5.0, + le=5.0 ) - value: torch.Tensor = Field( + value: torch.Tensor = pa.Field( dtype=torch.float32, shape=(None,), - checks=Check.greater_than(0.0) + gt=0.0 ) class Config: @@ -564,19 +565,19 @@ class ActorCriticOutput(pa.TensorDictModel): 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: @@ -874,7 +875,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/ @@ -886,6 +889,8 @@ pandera/ │ └── tensordict_engine.py # Engine for torch dtype registry ├── typing/ │ └── tensordict.py # Annotation types (TensorDict, TensorClass) +└── tensordict_api.py # Entry point: import pandera.tensordict as pa +``` └── tensordict.py # Entry point: import pandera.tensordict as pa ``` From b1e7fcbc04bd25a8ed641faee765b231b923c74d Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 21:01:18 -0400 Subject: [PATCH 02/27] test(tensordict): add unit tests for Phase 1 implementation - Add test_tensordict_container.py for Tensor component and TensorDictSchema - Add test_tensordict_model.py for TensorDictModel class-based schemas - Add test_tensordict_engine.py for tensordict_engine dtype resolution - Add __init__.py for test package --- tests/tensordict/__init__.py | 1 + tests/tensordict/test_tensordict_container.py | 358 ++++++++++++++++++ tests/tensordict/test_tensordict_engine.py | 109 ++++++ tests/tensordict/test_tensordict_model.py | 190 ++++++++++ 4 files changed, 658 insertions(+) create mode 100644 tests/tensordict/__init__.py create mode 100644 tests/tensordict/test_tensordict_container.py create mode 100644 tests/tensordict/test_tensordict_engine.py create mode 100644 tests/tensordict/test_tensordict_model.py 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..699d34730 --- /dev/null +++ b/tests/tensordict/test_tensordict_container.py @@ -0,0 +1,358 @@ +"""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_api import Tensor + + tensor = Tensor(dtype=torch.float32, shape=(None, 10)) + assert tensor.dtype == 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_api import Tensor + + tensor = Tensor( + dtype=torch.float32, + shape=(None, 10), + checks=Check.greater_than(0.0), + ) + assert tensor.dtype == torch.float32 + assert tensor.shape == (None, 10) + assert len(tensor.checks) == 1 + + def test_tensor_repr(self): + """Test Tensor repr.""" + from pandera.tensordict_api 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_api 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.columns + assert "action" in schema.columns + assert schema.batch_size == (32,) + + def test_tensordict_schema_from_list(self): + """Test TensorDictSchema creation from key list.""" + from pandera.tensordict_api import TensorDictSchema + + schema = TensorDictSchema(keys=["observation", "action"], batch_size=(32,)) + assert "observation" in schema.columns + assert "action" in schema.columns + + def test_tensordict_schema_repr(self): + """Test TensorDictSchema repr.""" + from pandera.tensordict_api 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_api 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_api 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.SchemaErrors): + schema.validate(td) + + def test_validate_missing_key(self): + """Test validation fails with missing key.""" + from pandera import errors + from pandera.tensordict_api 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.SchemaErrors): + schema.validate(td) + + def test_validate_wrong_dtype(self): + """Test validation fails with wrong dtype.""" + from pandera import errors + from pandera.tensordict_api 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.SchemaErrors): + schema.validate(td) + + def test_validate_wrong_shape(self): + """Test validation fails with wrong shape.""" + from pandera import errors + from pandera.tensordict_api 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.SchemaErrors): + schema.validate(td) + + def test_validate_lazy(self): + """Test lazy validation collects all errors.""" + from pandera import errors + from pandera.tensordict_api 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(16, 10), "action": torch.randn(32, 5)}, + batch_size=[16], + ) + + 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_api 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_api 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.SchemaErrors): + 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_api 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_api 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.SchemaErrors): + 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_api 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_api 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.SchemaErrors): + schema.validate(td) + + +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..5c3f6b1b2 --- /dev/null +++ b/tests/tensordict/test_tensordict_engine.py @@ -0,0 +1,109 @@ +"""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 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..1132523fb --- /dev/null +++ b/tests/tensordict/test_tensordict_model.py @@ -0,0 +1,190 @@ +"""Unit tests for TensorDictModel.""" + +import pytest + +try: + import torch + from tensordict import TensorDict +except ImportError: + torch = None + TensorDict = None + +torch_condition = pytest.mark.skipif(torch is None, reason="torch not installed") + + +@torch_condition +class TestTensorDictModel: + """Tests for TensorDictModel.""" + + def test_model_basic(self): + """Test basic TensorDictModel.""" + from pandera.tensordict_api import TensorDictModel + + class BasicModel(TensorDictModel): + observation: torch.Tensor + action: torch.Tensor + + schema = BasicModel.to_schema() + assert "observation" in schema.columns + assert "action" in schema.columns + + def test_model_with_field(self): + """Test TensorDictModel with Field.""" + from pandera.tensordict_api import Field, TensorDictModel + + class ModelWithField(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(None, 10)) + action: torch.Tensor = Field(dtype=torch.float32, shape=(None, 5)) + + schema = ModelWithField.to_schema() + assert "observation" in schema.columns + assert "action" in schema.columns + assert schema.columns["observation"].dtype == torch.float32 + + def test_model_with_config(self): + """Test TensorDictModel with Config.""" + from pandera.tensordict_api import TensorDictModel + + class ModelWithConfig(TensorDictModel): + observation: torch.Tensor + action: torch.Tensor + + class Config: + batch_size = (32,) + + schema = ModelWithConfig.to_schema() + assert schema.batch_size == (32,) + + def test_model_with_checks(self): + """Test TensorDictModel with check parameters.""" + from pandera.tensordict_api import Field, TensorDictModel + + class ModelWithChecks(TensorDictModel): + values: torch.Tensor = Field( + dtype=torch.float32, + shape=(None,), + gt=0.0, + lt=100.0, + ) + + schema = ModelWithChecks.to_schema() + assert "values" in schema.columns + assert len(schema.columns["values"].checks) == 2 + + def test_model_validate(self): + """Test TensorDictModel validate method.""" + from pandera.tensordict_api import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + action: torch.Tensor = Field(dtype=torch.float32, shape=(32, 5)) + + class Config: + batch_size = (32,) + + td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + batch_size=[32], + ) + + result = ModelForValidation.validate(td) + assert isinstance(result, TensorDict) + + def test_model_validate_invalid(self): + """Test TensorDictModel validate raises on invalid data.""" + from pandera import errors + from pandera.tensordict_api import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + action: torch.Tensor = Field(dtype=torch.float32, shape=(32, 5)) + + class Config: + batch_size = (32,) + + td = TensorDict( + {"observation": torch.randn(16, 10), "action": torch.randn(32, 5)}, + batch_size=[16], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td) + + def test_model_optional_field(self): + """Test TensorDictModel with optional field.""" + from pandera.tensordict_api import Field, TensorDictModel + + class ModelWithOptional(TensorDictModel): + observation: torch.Tensor + action: torch.Tensor | None + + schema = ModelWithOptional.to_schema() + assert "observation" in schema.columns + assert "action" in schema.columns + + def test_model_field_with_no_nan(self): + """Test TensorDictModel with no_nan check.""" + from pandera.tensordict_api import Field, TensorDictModel + + class ModelWithNoNan(TensorDictModel): + values: torch.Tensor = Field( + dtype=torch.float32, + shape=(None,), + no_nan=True, + ) + + schema = ModelWithNoNan.to_schema() + assert "values" in schema.columns + assert len(schema.columns["values"].checks) == 1 + + def test_model_field_with_isin(self): + """Test TensorDictModel with isin check.""" + from pandera.tensordict_api import Field, TensorDictModel + + class ModelWithIsin(TensorDictModel): + actions: torch.Tensor = Field( + dtype=torch.int64, + shape=(None,), + isin=[0, 1, 2, 3], + ) + + schema = ModelWithIsin.to_schema() + assert "actions" in schema.columns + assert len(schema.columns["actions"].checks) == 1 + + +@torch_condition +class TestTensorDictModelEdgeCases: + """Edge case tests for TensorDictModel.""" + + def test_model_missing_annotation(self): + """Test model raises error with missing annotation.""" + from pandera import errors + from pandera.tensordict_api import TensorDictModel + + class ModelWithMissingAnnotation(TensorDictModel): + observation = None # type: ignore + + with pytest.raises(errors.SchemaInitError): + ModelWithMissingAnnotation.to_schema() + + def test_model_empty(self): + """Test model with no fields.""" + from pandera.tensordict_api import TensorDictModel + + class EmptyModel(TensorDictModel): + pass + + schema = EmptyModel.to_schema() + assert schema.columns == {} + + def test_model_cannot_instantiate_directly(self): + """Test that TensorDictModel cannot be instantiated directly.""" + from pandera.tensordict_api import TensorDictModel + + with pytest.raises(NotImplementedError): + TensorDictModel() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 6888d1c63778b109812de5dfa8c8e60dae01d302 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 21:06:39 -0400 Subject: [PATCH 03/27] fix(tensordict): rename entry point and fix Check calling - Rename tensordict_api.py to tensordict.py per spec - Fix Check calling to use get_backend() pattern - Add tensorclass validation tests - Fix spec module layout duplicate entry - Update test imports --- pandera/backends/tensordict/base.py | 6 +- pandera/{tensordict_api.py => tensordict.py} | 0 specs/pytorch-tensordict.md | 5 -- tests/tensordict/test_tensordict_container.py | 67 ++++++++++++++----- tests/tensordict/test_tensordict_model.py | 24 +++---- 5 files changed, 66 insertions(+), 36 deletions(-) rename pandera/{tensordict_api.py => tensordict.py} (100%) diff --git a/pandera/backends/tensordict/base.py b/pandera/backends/tensordict/base.py index b8f1f1f5a..0bb9c5899 100644 --- a/pandera/backends/tensordict/base.py +++ b/pandera/backends/tensordict/base.py @@ -249,6 +249,8 @@ 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.columns.items(): if key not in check_obj: continue @@ -264,7 +266,9 @@ def _run_value_checks( for check_index, check in enumerate(tensor_schema.checks): try: - check_result = check(tensor_data, key) + check_result = check.get_backend(check_obj)(check)( + tensor_data, key + ) except Exception as exc: check_result = CoreCheckResult( passed=False, diff --git a/pandera/tensordict_api.py b/pandera/tensordict.py similarity index 100% rename from pandera/tensordict_api.py rename to pandera/tensordict.py diff --git a/specs/pytorch-tensordict.md b/specs/pytorch-tensordict.md index f3708b75e..4b6156bb6 100644 --- a/specs/pytorch-tensordict.md +++ b/specs/pytorch-tensordict.md @@ -887,10 +887,5 @@ pandera/ │ └── register.py # Backend registration ├── engines/ │ └── tensordict_engine.py # Engine for torch dtype registry -├── typing/ -│ └── tensordict.py # Annotation types (TensorDict, TensorClass) -└── tensordict_api.py # Entry point: import pandera.tensordict as pa -``` └── tensordict.py # Entry point: import pandera.tensordict as pa - ``` diff --git a/tests/tensordict/test_tensordict_container.py b/tests/tensordict/test_tensordict_container.py index 699d34730..53e8f11aa 100644 --- a/tests/tensordict/test_tensordict_container.py +++ b/tests/tensordict/test_tensordict_container.py @@ -19,7 +19,7 @@ class TestTensorComponent: def test_tensor_creation(self): """Test Tensor component creation.""" - from pandera.tensordict_api import Tensor + from pandera.tensordict import Tensor tensor = Tensor(dtype=torch.float32, shape=(None, 10)) assert tensor.dtype == torch.float32 @@ -28,7 +28,7 @@ def test_tensor_creation(self): def test_tensor_with_checks(self): """Test Tensor component with checks.""" from pandera import Check - from pandera.tensordict_api import Tensor + from pandera.tensordict import Tensor tensor = Tensor( dtype=torch.float32, @@ -41,7 +41,7 @@ def test_tensor_with_checks(self): def test_tensor_repr(self): """Test Tensor repr.""" - from pandera.tensordict_api import Tensor + from pandera.tensordict import Tensor tensor = Tensor(dtype=torch.float32, shape=(None, 10)) repr_str = repr(tensor) @@ -55,7 +55,7 @@ class TestTensorDictSchema: def test_tensordict_schema_creation(self): """Test TensorDictSchema creation.""" - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -70,7 +70,7 @@ def test_tensordict_schema_creation(self): def test_tensordict_schema_from_list(self): """Test TensorDictSchema creation from key list.""" - from pandera.tensordict_api import TensorDictSchema + from pandera.tensordict import TensorDictSchema schema = TensorDictSchema(keys=["observation", "action"], batch_size=(32,)) assert "observation" in schema.columns @@ -78,7 +78,7 @@ def test_tensordict_schema_from_list(self): def test_tensordict_schema_repr(self): """Test TensorDictSchema repr.""" - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -97,7 +97,7 @@ class TestTensorDictValidation: def test_validate_valid_tensordict(self): """Test validation of valid TensorDict.""" - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -121,7 +121,7 @@ def test_validate_valid_tensordict(self): def test_validate_invalid_batch_size(self): """Test validation fails with wrong batch size.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -141,7 +141,7 @@ def test_validate_invalid_batch_size(self): def test_validate_missing_key(self): """Test validation fails with missing key.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -162,7 +162,7 @@ def test_validate_missing_key(self): def test_validate_wrong_dtype(self): """Test validation fails with wrong dtype.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -182,7 +182,7 @@ def test_validate_wrong_dtype(self): def test_validate_wrong_shape(self): """Test validation fails with wrong shape.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -202,7 +202,7 @@ def test_validate_wrong_shape(self): def test_validate_lazy(self): """Test lazy validation collects all errors.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -230,7 +230,7 @@ class TestTensorDictSchemaChecks: def test_validate_with_value_check(self): """Test validation with value check.""" from pandera import Check - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -254,7 +254,7 @@ def test_validate_with_value_check(self): def test_validate_value_check_failure(self): """Test validation fails with value check failure.""" from pandera import Check, errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -282,7 +282,7 @@ class TestTensorDictSchemaBatchSize: def test_batch_size_with_none_dimension(self): """Test batch size with None dimension allows any size.""" - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -302,7 +302,7 @@ def test_batch_size_with_none_dimension(self): def test_batch_size_exact_match(self): """Test batch size must match exactly when specified.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -323,7 +323,7 @@ class TestShapeValidation: def test_shape_with_none_allows_any(self): """Test shape with None allows any size for that dimension.""" - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -339,7 +339,7 @@ def test_shape_with_none_allows_any(self): def test_shape_exact_match(self): """Test shape must match exactly when specified.""" from pandera import errors - from pandera.tensordict_api import Tensor, TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={ @@ -354,5 +354,36 @@ def test_shape_exact_match(self): 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 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index 1132523fb..7b05643f6 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -18,7 +18,7 @@ class TestTensorDictModel: def test_model_basic(self): """Test basic TensorDictModel.""" - from pandera.tensordict_api import TensorDictModel + from pandera.tensordict import TensorDictModel class BasicModel(TensorDictModel): observation: torch.Tensor @@ -30,7 +30,7 @@ class BasicModel(TensorDictModel): def test_model_with_field(self): """Test TensorDictModel with Field.""" - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelWithField(TensorDictModel): observation: torch.Tensor = Field(dtype=torch.float32, shape=(None, 10)) @@ -43,7 +43,7 @@ class ModelWithField(TensorDictModel): def test_model_with_config(self): """Test TensorDictModel with Config.""" - from pandera.tensordict_api import TensorDictModel + from pandera.tensordict import TensorDictModel class ModelWithConfig(TensorDictModel): observation: torch.Tensor @@ -57,7 +57,7 @@ class Config: def test_model_with_checks(self): """Test TensorDictModel with check parameters.""" - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelWithChecks(TensorDictModel): values: torch.Tensor = Field( @@ -73,7 +73,7 @@ class ModelWithChecks(TensorDictModel): def test_model_validate(self): """Test TensorDictModel validate method.""" - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelForValidation(TensorDictModel): observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) @@ -93,7 +93,7 @@ class Config: def test_model_validate_invalid(self): """Test TensorDictModel validate raises on invalid data.""" from pandera import errors - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelForValidation(TensorDictModel): observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) @@ -112,7 +112,7 @@ class Config: def test_model_optional_field(self): """Test TensorDictModel with optional field.""" - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelWithOptional(TensorDictModel): observation: torch.Tensor @@ -124,7 +124,7 @@ class ModelWithOptional(TensorDictModel): def test_model_field_with_no_nan(self): """Test TensorDictModel with no_nan check.""" - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelWithNoNan(TensorDictModel): values: torch.Tensor = Field( @@ -139,7 +139,7 @@ class ModelWithNoNan(TensorDictModel): def test_model_field_with_isin(self): """Test TensorDictModel with isin check.""" - from pandera.tensordict_api import Field, TensorDictModel + from pandera.tensordict import Field, TensorDictModel class ModelWithIsin(TensorDictModel): actions: torch.Tensor = Field( @@ -160,7 +160,7 @@ class TestTensorDictModelEdgeCases: def test_model_missing_annotation(self): """Test model raises error with missing annotation.""" from pandera import errors - from pandera.tensordict_api import TensorDictModel + from pandera.tensordict import TensorDictModel class ModelWithMissingAnnotation(TensorDictModel): observation = None # type: ignore @@ -170,7 +170,7 @@ class ModelWithMissingAnnotation(TensorDictModel): def test_model_empty(self): """Test model with no fields.""" - from pandera.tensordict_api import TensorDictModel + from pandera.tensordict import TensorDictModel class EmptyModel(TensorDictModel): pass @@ -180,7 +180,7 @@ class EmptyModel(TensorDictModel): def test_model_cannot_instantiate_directly(self): """Test that TensorDictModel cannot be instantiated directly.""" - from pandera.tensordict_api import TensorDictModel + from pandera.tensordict import TensorDictModel with pytest.raises(NotImplementedError): TensorDictModel() From 55888f133379dea724d9615d13bcfc275487da2e Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 21:23:37 -0400 Subject: [PATCH 04/27] test(tensordict): add comprehensive error case tests - Add TestTensorDictErrorCases with 14+ error validation tests - Add TestTensorDictModelErrorCases with model error tests - Add TestTensorDictEngineErrorCases with engine error tests - Test batch_size, dtype, shape, missing keys, wrong input types - Test error messages contain relevant info - Test lazy validation collects multiple errors --- tests/tensordict/test_tensordict_container.py | 257 ++++++++++++++++++ tests/tensordict/test_tensordict_engine.py | 38 +++ tests/tensordict/test_tensordict_model.py | 118 ++++++++ 3 files changed, 413 insertions(+) diff --git a/tests/tensordict/test_tensordict_container.py b/tests/tensordict/test_tensordict_container.py index 53e8f11aa..65d8f72a2 100644 --- a/tests/tensordict/test_tensordict_container.py +++ b/tests/tensordict/test_tensordict_container.py @@ -385,5 +385,262 @@ class TCData: 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 TensorDictSchema + + schema = TensorDictSchema(keys={"x": Tensor(dtype=torch.float32)}) + + with pytest.raises(TypeError, match="Expected TensorDict"): + 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.SchemaErrors): + 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.SchemaErrors): + 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.SchemaErrors): + 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.SchemaErrors): + 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.SchemaErrors): + 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.SchemaErrors): + 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.SchemaErrors): + 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(32, 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.SchemaErrors as e: + error_messages = [err.message for err in e.schema_errors] + 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.SchemaErrors as e: + error_messages = [err.message for err in e.schema_errors] + 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.SchemaErrors as e: + error_messages = [err.message for err in e.schema_errors] + 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.SchemaErrors as e: + error_messages = [err.message for err in e.schema_errors] + 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.tensordict import Tensor, TensorDictSchema + from pandera.errors import SchemaErrorReason + + 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.SchemaErrors as e: + reason_codes = [err.reason_code for err in e.schema_errors] + assert SchemaErrorReason.WRONG_TYPE 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 index 5c3f6b1b2..84e6aa272 100644 --- a/tests/tensordict/test_tensordict_engine.py +++ b/tests/tensordict/test_tensordict_engine.py @@ -92,6 +92,44 @@ def test_engine_multiple_dtypes(self): 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.engines import tensordict_engine + from pandera import errors + + 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.""" diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index 7b05643f6..adbac13ed 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -186,5 +186,123 @@ def test_model_cannot_instantiate_directly(self): TensorDictModel() +@torch_condition +class TestTensorDictModelErrorCases: + """Comprehensive error case tests for TensorDictModel.""" + + def test_model_validate_wrong_batch_size(self): + """Test model validate fails with wrong batch size.""" + from pandera import errors + from pandera.tensordict import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + + class Config: + batch_size = (32,) + + td = TensorDict( + {"observation": torch.randn(16, 10)}, + batch_size=[16], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td) + + def test_model_validate_wrong_dtype(self): + """Test model validate fails with wrong dtype.""" + from pandera import errors + from pandera.tensordict import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + + class Config: + batch_size = (32,) + + td = TensorDict( + {"observation": torch.randn(32, 10).to(torch.int32)}, + batch_size=[32], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td) + + def test_model_validate_missing_key(self): + """Test model validate fails with missing key.""" + from pandera import errors + from pandera.tensordict import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + action: torch.Tensor = Field(dtype=torch.float32, shape=(32, 5)) + + class Config: + batch_size = (32,) + + td = TensorDict( + {"observation": torch.randn(32, 10)}, + batch_size=[32], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td) + + def test_model_validate_wrong_shape(self): + """Test model validate fails with wrong shape.""" + from pandera import errors + from pandera.tensordict import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + + class Config: + batch_size = (32,) + + td = TensorDict( + {"observation": torch.randn(32, 20)}, + batch_size=[32], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td) + + def test_model_validate_value_check_failure(self): + """Test model validate fails with value check failure.""" + from pandera import errors + from pandera.tensordict import Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + values: torch.Tensor = Field( + dtype=torch.float32, + shape=(10,), + gt=0.0, + ) + + class Config: + 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.SchemaErrors): + ModelForValidation.validate(td) + + def test_model_config_batch_size(self): + """Test model config batch_size is used.""" + from pandera.tensordict import Field, TensorDictModel + + class ModelWithConfig(TensorDictModel): + observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + + class Config: + batch_size = (32,) + + schema = ModelWithConfig.to_schema() + assert schema.batch_size == (32,) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From fe466962bd2d5ce4608d68357a96e3e50fd83340 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 21:52:28 -0400 Subject: [PATCH 05/27] fix(tensordict): address dtype, SchemaError, and validation issues - tensordict_engine: strip 'torch.' prefix from string dtypes - base.py: use data= not schema_context= for SchemaError - base.py: implement failure_cases_metadata for lazy validation - base.py: fix SchemaErrorReason codes (COLUMN_NOT_IN_DATAFRAME, CHECK_ERROR) - base.py: fix empty check result message - tests: fix dtype comparisons, SchemaError vs SchemaErrors, error messages --- pandera/api/tensordict/components.py | 19 ++++- pandera/api/tensordict/types.py | 15 ++-- pandera/backends/tensordict/base.py | 77 +++++++++++-------- pandera/backends/tensordict/register.py | 4 +- pandera/engines/tensordict_engine.py | 31 ++++++-- tests/tensordict/test_tensordict_container.py | 68 ++++++++-------- 6 files changed, 134 insertions(+), 80 deletions(-) diff --git a/pandera/api/tensordict/components.py b/pandera/api/tensordict/components.py index 48b518830..60c538ffd 100644 --- a/pandera/api/tensordict/components.py +++ b/pandera/api/tensordict/components.py @@ -46,6 +46,10 @@ def __init__( :raises SchemaInitError: if impossible to build schema from parameters """ + if dtype is not None and torch is not None: + from pandera.engines import tensordict_engine + dtype = tensordict_engine.Engine.dtype(dtype) if dtype else None + super().__init__( dtype=dtype, checks=checks, @@ -59,8 +63,21 @@ def __init__( ) self.shape = shape + @property + def dtype(self): + return self._dtype + + @dtype.setter + def dtype(self, value) -> None: + if torch is None: + self._dtype = value + else: + from pandera.engines import tensordict_engine + self._dtype = tensordict_engine.Engine.dtype(value) if value else None + def __repr__(self) -> str: + dtype_str = self._dtype.type if hasattr(self._dtype, 'type') else str(self._dtype) return ( - f"Tensor(dtype={self.dtype}, shape={self.shape}, " + f"Tensor(dtype={dtype_str}, shape={self.shape}, " f"nullable={self.nullable})" ) diff --git a/pandera/api/tensordict/types.py b/pandera/api/tensordict/types.py index 1421c7457..8ffc50173 100644 --- a/pandera/api/tensordict/types.py +++ b/pandera/api/tensordict/types.py @@ -1,19 +1,22 @@ """Type definitions for pandera tensordict integration.""" -from typing import TYPE_CING, Any, Union - -import torch +from typing import TYPE_CHECKING, Any, Union from pandera.api.checks import Check -if TYPE_CING: +try: + import torch +except ImportError: + torch = None + +if TYPE_CHECKING: from tensordict import TensorDict, tensorclass TensorDictCheckObjects = Any -TensorDictInputType = Union["TensorDict", "tensorclass", torch.Tensor] +TensorDictInputType = Union["TensorDict", "tensorclass", torch.Tensor] # type: ignore[union-attr] -TensorDtypeInputTypes = Union[torch.dtype, str, type] +TensorDtypeInputTypes = Union[torch.dtype, str, type] # type: ignore[union-attr] def is_tensordict(obj: Any) -> bool: diff --git a/pandera/backends/tensordict/base.py b/pandera/backends/tensordict/base.py index 0bb9c5899..baa185dad 100644 --- a/pandera/backends/tensordict/base.py +++ b/pandera/backends/tensordict/base.py @@ -103,9 +103,9 @@ def _check_batch_size( ) error = SchemaError( schema=schema, - schema_context={"batch_size": actual_batch_size}, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.WRONG_TYPE, + reason_code=SchemaErrorReason.WRONG_DATATYPE, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -123,9 +123,9 @@ def _check_batch_size( ) error = SchemaError( schema=schema, - schema_context={"batch_size": actual_batch_size}, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.WRONG_TYPE, + reason_code=SchemaErrorReason.WRONG_DATATYPE, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -140,9 +140,9 @@ def _check_keys(self, check_obj, schema, error_handler, lazy): error_msg = f"Missing key '{key}' in TensorDict" error = SchemaError( schema=schema, - schema_context={"keys": list(check_obj.keys())}, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.MISSING_COLUMN, + reason_code=SchemaErrorReason.COLUMN_NOT_IN_DATAFRAME, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -164,9 +164,9 @@ def _check_dtypes_and_shapes( error_msg = f"Key '{key}' is not a torch.Tensor" error = SchemaError( schema=schema, - schema_context={"key": key, "type": type(tensor)}, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.WRONG_TYPE, + reason_code=SchemaErrorReason.WRONG_DATATYPE, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -177,19 +177,17 @@ def _check_dtypes_and_shapes( if tensor_schema.dtype is not None: expected_dtype = tensor_schema.dtype.type - if tensor.dtype != expected_dtype: + actual_dtype = tensor.dtype + if actual_dtype != expected_dtype: error_msg = ( f"Key '{key}': expected dtype {expected_dtype}, " - f"got {tensor.dtype}" + f"got {actual_dtype}" ) error = SchemaError( schema=schema, - schema_context={ - "key": key, - "dtype": str(tensor.dtype), - }, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.WRONG_TYPE, + reason_code=SchemaErrorReason.WRONG_DATATYPE, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -208,12 +206,9 @@ def _check_dtypes_and_shapes( ) error = SchemaError( schema=schema, - schema_context={ - "key": key, - "shape": actual_shape, - }, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.WRONG_TYPE, + reason_code=SchemaErrorReason.WRONG_DATATYPE, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -232,12 +227,9 @@ def _check_dtypes_and_shapes( ) error = SchemaError( schema=schema, - schema_context={ - "key": key, - "shape": actual_shape, - }, + data=check_obj, message=error_msg, - reason_code=SchemaErrorReason.WRONG_TYPE, + reason_code=SchemaErrorReason.WRONG_DATATYPE, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -274,22 +266,23 @@ def _run_value_checks( passed=False, check=check, check_index=check_index, - reason_code=SchemaErrorReason.CHECK_FAILED, + reason_code=SchemaErrorReason.CHECK_ERROR, message=str(exc), ) if not check_result.passed: + check_msg = check_result.message or "check failed" error_msg = ( f"Check '{check}' failed for key '{key}': " - f"{check_result.message}" + f"{check_msg}" ) error = SchemaError( schema=schema, - schema_context={"key": key}, + data=check_obj, message=error_msg, check=check, check_index=check_index, - reason_code=SchemaErrorReason.CHECK_FAILED, + reason_code=SchemaErrorReason.CHECK_ERROR, ) error_handler.collect_error( get_error_category(error.reason_code), @@ -335,7 +328,31 @@ def failure_cases_metadata( self, schema_name: str, schema_errors: list[SchemaError] ): """Get failure cases metadata for lazy validation.""" - return None + from collections import defaultdict + + from pandera.errors import FailureCaseMetadata + + from pandera.api.base.error_handler import ErrorHandler + + 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 = 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.""" diff --git a/pandera/backends/tensordict/register.py b/pandera/backends/tensordict/register.py index 415263855..814113d40 100644 --- a/pandera/backends/tensordict/register.py +++ b/pandera/backends/tensordict/register.py @@ -1,9 +1,9 @@ """Register TensorDict backends.""" from functools import lru_cache -from typing import TYPE_CING, Any +from typing import TYPE_CHECKING, Any -if TYPE_CING: +if TYPE_CHECKING: from tensordict import TensorDict, tensorclass diff --git a/pandera/engines/tensordict_engine.py b/pandera/engines/tensordict_engine.py index f968444b2..9bfe6a51c 100644 --- a/pandera/engines/tensordict_engine.py +++ b/pandera/engines/tensordict_engine.py @@ -25,10 +25,23 @@ class DataType(dtypes.DataType): def __init__(self, dtype: Any): super().__init__() - object.__setattr__(self, "type", torch.dtype(dtype)) + 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): - object.__setattr__(self, "type", torch.dtype(self.type)) + 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.""" @@ -63,11 +76,15 @@ def dtype(cls, data_type: Any) -> dtypes.DataType: """Convert input into a PyTorch-compatible Pandera DataType.""" try: return engine.Engine.dtype(cls, data_type) - except TypeError: + except (TypeError, ValueError): try: - torch_dtype = torch.dtype(data_type) - return DataType(torch_dtype) - except TypeError: + 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'." @@ -113,7 +130,7 @@ def _register_torch_dtypes(): try: torch_dtype = getattr(torch, dtype_name, None) if torch_dtype is not None: - Engine.register_dtype(DataType(torch_dtype)) + Engine.register_dtype(DataType(dtype_name)) except (AttributeError, ValueError): pass diff --git a/tests/tensordict/test_tensordict_container.py b/tests/tensordict/test_tensordict_container.py index 65d8f72a2..1bc144f4b 100644 --- a/tests/tensordict/test_tensordict_container.py +++ b/tests/tensordict/test_tensordict_container.py @@ -22,7 +22,7 @@ def test_tensor_creation(self): from pandera.tensordict import Tensor tensor = Tensor(dtype=torch.float32, shape=(None, 10)) - assert tensor.dtype == torch.float32 + assert tensor.dtype.type == torch.float32 assert tensor.shape == (None, 10) def test_tensor_with_checks(self): @@ -35,7 +35,7 @@ def test_tensor_with_checks(self): shape=(None, 10), checks=Check.greater_than(0.0), ) - assert tensor.dtype == torch.float32 + assert tensor.dtype.type == torch.float32 assert tensor.shape == (None, 10) assert len(tensor.checks) == 1 @@ -135,7 +135,7 @@ def test_validate_invalid_batch_size(self): batch_size=[16], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_validate_missing_key(self): @@ -156,7 +156,7 @@ def test_validate_missing_key(self): batch_size=[32], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_validate_wrong_dtype(self): @@ -176,7 +176,7 @@ def test_validate_wrong_dtype(self): batch_size=[32], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_validate_wrong_shape(self): @@ -196,7 +196,7 @@ def test_validate_wrong_shape(self): batch_size=[32], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_validate_lazy(self): @@ -206,15 +206,15 @@ def test_validate_lazy(self): schema = TensorDictSchema( keys={ - "observation": Tensor(dtype=torch.float32, shape=(32, 10)), - "action": Tensor(dtype=torch.float32, shape=(32, 5)), + "observation": Tensor(dtype=torch.float32, shape=(None, 10)), + "action": Tensor(dtype=torch.float32, shape=(None, 5)), }, batch_size=(32,), ) td = TensorDict( - {"observation": torch.randn(16, 10), "action": torch.randn(32, 5)}, - batch_size=[16], + {"observation": torch.randn(32, 10), "action": torch.randn(16, 5)}, + batch_size=None, ) with pytest.raises(errors.SchemaErrors) as exc_info: @@ -272,7 +272,7 @@ def test_validate_value_check_failure(self): batch_size=[10], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) @@ -313,7 +313,7 @@ def test_batch_size_exact_match(self): td = TensorDict({"observation": torch.randn(64, 10)}, batch_size=[64]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) @@ -350,7 +350,7 @@ def test_shape_exact_match(self): td = TensorDict({"data": torch.randn(32, 20)}, batch_size=[32]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) @@ -392,11 +392,11 @@ class TestTensorDictErrorCases: def test_invalid_input_type(self): """Test validation fails with non-TensorDict input.""" from pandera import errors - from pandera.tensordict import TensorDictSchema + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema(keys={"x": Tensor(dtype=torch.float32)}) - with pytest.raises(TypeError, match="Expected TensorDict"): + with pytest.raises(errors.BackendNotFoundError, match="Backend not found"): schema.validate({"x": torch.randn(10)}) def test_invalid_tensor_in_tensordict(self): @@ -411,7 +411,7 @@ def test_invalid_tensor_in_tensordict(self): td = TensorDict({"x": "not a tensor"}, batch_size=[10]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_batch_size_length_mismatch(self): @@ -426,7 +426,7 @@ def test_batch_size_length_mismatch(self): td = TensorDict({"x": torch.randn(10)}, batch_size=[10]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_batch_size_value_mismatch_first_dim(self): @@ -441,7 +441,7 @@ def test_batch_size_value_mismatch_first_dim(self): td = TensorDict({"x": torch.randn(5)}, batch_size=[5]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_batch_size_value_mismatch_second_dim(self): @@ -456,7 +456,7 @@ def test_batch_size_value_mismatch_second_dim(self): td = TensorDict({"x": torch.randn(10, 3)}, batch_size=[10, 3]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_multiple_missing_keys(self): @@ -475,7 +475,7 @@ def test_multiple_missing_keys(self): td = TensorDict({"a": torch.randn(10)}, batch_size=[10]) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_multiple_dtype_mismatches(self): @@ -496,7 +496,7 @@ def test_multiple_dtype_mismatches(self): batch_size=[10], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_multiple_shape_errors(self): @@ -517,7 +517,7 @@ def test_multiple_shape_errors(self): batch_size=[10], ) - with pytest.raises(errors.SchemaErrors): + with pytest.raises(errors.SchemaError): schema.validate(td) def test_lazy_collects_all_errors(self): @@ -537,7 +537,7 @@ def test_lazy_collects_all_errors(self): td = TensorDict( { "a": torch.randn(16, 5), - "b": torch.randn(32, 8), + "b": torch.randn(16, 8), }, batch_size=[16], ) @@ -561,8 +561,8 @@ def test_error_message_content_batch_size(self): try: schema.validate(td) - except errors.SchemaErrors as e: - error_messages = [err.message for err in e.schema_errors] + 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): @@ -579,8 +579,8 @@ def test_error_message_content_dtype(self): try: schema.validate(td) - except errors.SchemaErrors as e: - error_messages = [err.message for err in e.schema_errors] + 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): @@ -597,8 +597,8 @@ def test_error_message_content_shape(self): try: schema.validate(td) - except errors.SchemaErrors as e: - error_messages = [err.message for err in e.schema_errors] + 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): @@ -618,8 +618,8 @@ def test_error_message_content_missing_key(self): try: schema.validate(td) - except errors.SchemaErrors as e: - error_messages = [err.message for err in e.schema_errors] + 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): @@ -637,9 +637,9 @@ def test_schema_error_reason_codes(self): try: schema.validate(td) - except errors.SchemaErrors as e: - reason_codes = [err.reason_code for err in e.schema_errors] - assert SchemaErrorReason.WRONG_TYPE in reason_codes + except errors.SchemaError as e: + reason_codes = [e.reason_code] + assert SchemaErrorReason.WRONG_DATATYPE in reason_codes if __name__ == "__main__": From dc1b99a02b5610cf9fff13c1719ea71eaa5f463f Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 21:53:20 -0400 Subject: [PATCH 06/27] fix(tensordict_engine): fix coerce_value to use torch.tensor --- pandera/engines/tensordict_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandera/engines/tensordict_engine.py b/pandera/engines/tensordict_engine.py index 9bfe6a51c..8d34f9552 100644 --- a/pandera/engines/tensordict_engine.py +++ b/pandera/engines/tensordict_engine.py @@ -49,7 +49,7 @@ def coerce(self, data_container: torch.Tensor) -> torch.Tensor: def coerce_value(self, value: Any) -> Any: """Coerce a value to the particular type.""" - return self.type(value) + return torch.tensor(value, dtype=self.type) def try_coerce(self, data_container: torch.Tensor) -> torch.Tensor: try: From 66773f2ade2bd7936dcf10db318a365a4ad1fa5e Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 22:10:18 -0400 Subject: [PATCH 07/27] docs: add PyTorch TensorDict data validation guide - Add pytorch_guide/ with index, schema, model, checks, and error reporting docs - Add reference/pytorch.rst API reference - Update index.md and reference/index.md to include pytorch guide --- docs/source/index.md | 1 + docs/source/pytorch_guide/error_reporting.md | 64 ++++++++++ docs/source/pytorch_guide/index.md | 88 ++++++++++++++ .../source/pytorch_guide/tensordict_checks.md | 73 ++++++++++++ docs/source/pytorch_guide/tensordict_model.md | 70 +++++++++++ .../source/pytorch_guide/tensordict_schema.md | 91 ++++++++++++++ docs/source/reference/index.md | 3 + docs/source/reference/pytorch.rst | 111 ++++++++++++++++++ 8 files changed, 501 insertions(+) create mode 100644 docs/source/pytorch_guide/error_reporting.md create mode 100644 docs/source/pytorch_guide/index.md create mode 100644 docs/source/pytorch_guide/tensordict_checks.md create mode 100644 docs/source/pytorch_guide/tensordict_model.md create mode 100644 docs/source/pytorch_guide/tensordict_schema.md create mode 100644 docs/source/reference/pytorch.rst 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..4a30045a5 --- /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" - {err.message}") + 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 \ No newline at end of file diff --git a/docs/source/pytorch_guide/index.md b/docs/source/pytorch_guide/index.md new file mode 100644 index 000000000..76ba56f67 --- /dev/null +++ b/docs/source/pytorch_guide/index.md @@ -0,0 +1,88 @@ +--- +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): + observation: torch.Tensor + action: torch.Tensor + + @classmethod + def _field(cls, field): + return field(dtype=torch.float32, shape=(None, 10)) + +# Validate using the model +td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + batch_size=[32], +) +RL.validate(td) +``` + +## 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 \ No newline at end of file diff --git a/docs/source/pytorch_guide/tensordict_checks.md b/docs/source/pytorch_guide/tensordict_checks.md new file mode 100644 index 000000000..e472ada39 --- /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, 1.0])}, + 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 \ No newline at end of file diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md new file mode 100644 index 000000000..dc4399389 --- /dev/null +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -0,0 +1,70 @@ +--- +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): + observation: torch.Tensor + action: torch.Tensor + + @classmethod + def _field(cls, field): + return field(dtype=torch.float32, shape=(None, 10)) +``` + +The `_field()` classmethod customizes the default field configuration. + +## Validate with a model + +```{code-cell} python +from tensordict import TensorDict + +td = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + batch_size=[32], +) +validated = RL.validate(td) +``` + +## Model configuration + +Use {class}`~pandera.api.tensordict.model_components.Field` with the `_field()` method: + +- `dtype`: Torch dtype +- `shape`: Expected shape tuple +- `checks`: List of Check instances +- `nullable`: Whether the key can be missing +- `default`: Default value if missing + +```{code-cell} python +class RLWithConfig(pa.TensorDictModel): + observation: torch.Tensor + action: torch.Tensor + reward: torch.Tensor + + @classmethod + def _field(cls, field): + return field( + dtype=torch.float32, + shape=(None,), + checks=[Check.greater_than(0.0)], + ) +``` + +## See also + +- {ref}`pytorch-tensordict-schema` — dictionary-based schema +- {ref}`pytorch-checks` — value checks +- {ref}`configuration` — validation configuration \ No newline at end of file diff --git a/docs/source/pytorch_guide/tensordict_schema.md b/docs/source/pytorch_guide/tensordict_schema.md new file mode 100644 index 000000000..9e74214aa --- /dev/null +++ b/docs/source/pytorch_guide/tensordict_schema.md @@ -0,0 +1,91 @@ +--- +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: + +```{code-cell} python +td_invalid = TensorDict( + {"observation": torch.randn(16, 10), "action": torch.randn(32, 5)}, + batch_size=[16], +) + +try: + schema.validate(td_invalid, lazy=True) +except pa.errors.SchemaErrors as exc: + print(exc) +``` + +## 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 \ No newline at end of file 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..0ca1865dd --- /dev/null +++ b/docs/source/reference/pytorch.rst @@ -0,0 +1,111 @@ +.. _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 + +Typing +------ + +.. autosummary:: + :toctree: generated + :template: class.rst + :nosignatures: + + pandera.typing.tensordict.TensorDict + pandera.typing.tensordict.Tensorclass + +Abstract base classes +--------------------- + +.. autosummary:: + :toctree: generated + :template: class.rst + :nosignatures: + + pandera.api.tensordict.base.BaseTensorDictSchema + +Check object types +------------------ + +Types accepted by TensorDict :class:`~pandera.api.checks.Check` backends. + +.. autosummary:: + :toctree: generated + :nosignatures: + + pandera.api.tensordict.types.TensorDictData + pandera.api.tensordict.types.TENSORDICT_CHECK_OBJECT_TYPES + +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 \ No newline at end of file From ae8bc77b676605070775f723f105807d9856f96c Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 13 Apr 2026 23:12:11 -0400 Subject: [PATCH 08/27] feat(torch): add torch extra and update config files - Add torch/tensordict extra to pyproject.toml - Add torch to DATAFRAME_EXTRAS in noxfile.py - Fix tensordict/types.py to handle missing torch gracefully - Add torch and tensordict to requirements.txt and environment.yml - Remove non-existent modules from pytorch reference --- docs/source/reference/pytorch.rst | 33 ------------------------------- environment.yml | 4 ++++ noxfile.py | 17 ++++++++++++++++ pandera/api/tensordict/types.py | 17 ++++++++++------ pyproject.toml | 6 ++++++ requirements.txt | 2 ++ 6 files changed, 40 insertions(+), 39 deletions(-) diff --git a/docs/source/reference/pytorch.rst b/docs/source/reference/pytorch.rst index 0ca1865dd..240d389c1 100644 --- a/docs/source/reference/pytorch.rst +++ b/docs/source/reference/pytorch.rst @@ -50,39 +50,6 @@ Declarative models pandera.api.tensordict.model_components.Field -Typing ------- - -.. autosummary:: - :toctree: generated - :template: class.rst - :nosignatures: - - pandera.typing.tensordict.TensorDict - pandera.typing.tensordict.Tensorclass - -Abstract base classes ---------------------- - -.. autosummary:: - :toctree: generated - :template: class.rst - :nosignatures: - - pandera.api.tensordict.base.BaseTensorDictSchema - -Check object types ------------------- - -Types accepted by TensorDict :class:`~pandera.api.checks.Check` backends. - -.. autosummary:: - :toctree: generated - :nosignatures: - - pandera.api.tensordict.types.TensorDictData - pandera.api.tensordict.types.TENSORDICT_CHECK_OBJECT_TYPES - Configuration ------------- 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..129a817a9 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 @@ -165,6 +177,10 @@ def _testing_requirements( req = "ibis-framework[duckdb] >= 11.0.0" if req == "polars": req = f"polars=={polars}" + if req == "torch": + req = "torch" + if req == "tensordict": + req = "tensordict" # for some reason uv will try to install an old version of dask, # have to specifically pin dask[dataframe] to a higher version @@ -199,6 +215,7 @@ def _testing_requirements( "dask", "ibis", "xarray", + "torch", } for extra in OPTIONAL_DEPENDENCIES: if extra == "pandas": diff --git a/pandera/api/tensordict/types.py b/pandera/api/tensordict/types.py index 8ffc50173..f5a3bfb3b 100644 --- a/pandera/api/tensordict/types.py +++ b/pandera/api/tensordict/types.py @@ -9,14 +9,19 @@ except ImportError: torch = None -if TYPE_CHECKING: - from tensordict import TensorDict, tensorclass +if torch is not None: + if TYPE_CHECKING: + from tensordict import TensorDict, tensorclass -TensorDictCheckObjects = Any + TensorDictCheckObjects = Any -TensorDictInputType = Union["TensorDict", "tensorclass", torch.Tensor] # type: ignore[union-attr] + TensorDictInputType = Union["TensorDict", "tensorclass", torch.Tensor] # type: ignore[union-attr] -TensorDtypeInputTypes = Union[torch.dtype, str, type] # type: ignore[union-attr] + TensorDtypeInputTypes = Union[torch.dtype, str, type] # type: ignore[union-attr] +else: + TensorDictCheckObjects = Any + TensorDictInputType = Any + TensorDtypeInputTypes = Any def is_tensordict(obj: Any) -> bool: @@ -26,4 +31,4 @@ def is_tensordict(obj: Any) -> bool: return isinstance(obj, (TensorDict, tensorclass)) except ImportError: - return False + return False \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0ced750a6..1467167a4 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] 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 From 2fda28db05a170ce9308c77470854b07de420e13 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 14 Apr 2026 10:23:30 -0400 Subject: [PATCH 09/27] Fix TensorDictModel implementation to use Field function and DataType type annotation - Expose errors module and DataType in pandera.tensordict namespace - Use DataType (from tensordict_engine) instead of torch.Tensor for type annotations - Replace _field classmethod with Field function for field configuration - Fix validate to be a classmethod for correct class-level calling - Update tests and documentation to match other pandera backends --- docs/source/pytorch_guide/tensordict_model.md | 34 ++-- pandera/api/tensordict/model.py | 15 +- pandera/api/tensordict/model_components.py | 30 ++-- pandera/engines/tensordict_engine.py | 2 + pandera/tensordict.py | 8 +- tests/tensordict/test_tensordict_model.py | 146 +++++++++++------- 6 files changed, 139 insertions(+), 96 deletions(-) diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md index dc4399389..d7ef78b24 100644 --- a/docs/source/pytorch_guide/tensordict_model.md +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -14,17 +14,15 @@ schema that maps directly to a {class}`~tensordict.TensorDict`. ```{code-cell} python import torch import pandera.tensordict as pa +from pandera import Check class RL(pa.TensorDictModel): - observation: torch.Tensor - action: torch.Tensor - - @classmethod - def _field(cls, field): - return field(dtype=torch.float32, shape=(None, 10)) + observation: pa.DataType + action: pa.DataType ``` -The `_field()` classmethod customizes the default field configuration. +Use {class}`~pandera.tensordict.DataType` as the type annotation and +{func}`~pandera.tensordict.Field` to customize field options. ## Validate with a model @@ -38,9 +36,9 @@ td = TensorDict( validated = RL.validate(td) ``` -## Model configuration +## Field configuration -Use {class}`~pandera.api.tensordict.model_components.Field` with the `_field()` method: +Use {func}`~pandera.tensordict.Field` to customize field options: - `dtype`: Torch dtype - `shape`: Expected shape tuple @@ -50,17 +48,13 @@ Use {class}`~pandera.api.tensordict.model_components.Field` with the `_field()` ```{code-cell} python class RLWithConfig(pa.TensorDictModel): - observation: torch.Tensor - action: torch.Tensor - reward: torch.Tensor - - @classmethod - def _field(cls, field): - return field( - dtype=torch.float32, - shape=(None,), - checks=[Check.greater_than(0.0)], - ) + observation: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 10)) + action: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 5)) + reward: pa.DataType = pa.Field( + dtype=torch.float32, + shape=(None,), + checks=[Check.greater_than(0.0)], + ) ``` ## See also diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index 470bb959c..b37022585 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -50,10 +50,10 @@ def build_schema_(cls, **kwargs) -> TensorDictSchema: keys=columns, checks=cls.__root_checks__, batch_size=cls.__config__.batch_size, - name=cls.__config__.name, - title=cls.__config__.title, - description=cls.__config__.description, - coerce=cls.__config__.coerce, + name=getattr(cls.__config__, "name", None), + title=getattr(cls.__config__, "title", None), + description=getattr(cls.__config__, "description", None), + coerce=getattr(cls.__config__, "coerce", False), **kwargs, ) @@ -65,7 +65,7 @@ def _build_columns( ) -> dict[str, Tensor]: columns = {} for name, (annotation, field) in fields.items(): - dtype = annotation.arg + dtype = annotation.raw_annotation if dtype is None: raise SchemaInitError( f"expected annotation for field '{name}'" @@ -156,8 +156,9 @@ def __init_subclass__(cls, **kwargs): def to_schema(cls, **kwargs) -> TensorDictSchema: return cls.build_schema_(**kwargs) - def validate(self, check_obj, *args, **kwargs): - schema = self.to_schema() + @classmethod + def validate(cls, check_obj, *args, **kwargs): + schema = cls.to_schema() return schema.validate(check_obj, *args, **kwargs) def __new__(cls, *args, **kwargs) -> Any: diff --git a/pandera/api/tensordict/model_components.py b/pandera/api/tensordict/model_components.py index 68f28d17e..8e943330f 100644 --- a/pandera/api/tensordict/model_components.py +++ b/pandera/api/tensordict/model_components.py @@ -3,6 +3,7 @@ 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 @@ -58,24 +59,24 @@ def to_tensor_kwargs( checks: Any = None, ) -> dict[str, Any]: """Keyword args for :class:`~pandera.api.tensordict.components.Tensor`.""" - req = self.required and not optional - return self._get_schema_properties( - dtype, - checks=checks, - required=req, - shape=self.shape, - alias=self.alias, - title=self.title, - description=self.description, - default=self.default, - metadata=self.metadata, - nullable=self.nullable, - coerce=self.coerce, - ) + if self.dtype_kwargs: + dtype = dtype(**self.dtype_kwargs) + return { + "dtype": dtype, + "checks": to_checklist(self.checks) + to_checklist(checks), + "shape": self.shape, + "name": self.alias, + "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, @@ -109,7 +110,6 @@ def Field( default: Any | None = None, metadata: dict[str, Any] | None = None, required: bool = True, - shape: tuple[int | None, ...] | None = None, **kwargs: Any, ) -> Any: """Field specification for TensorDict models.""" diff --git a/pandera/engines/tensordict_engine.py b/pandera/engines/tensordict_engine.py index 8d34f9552..023bd22bd 100644 --- a/pandera/engines/tensordict_engine.py +++ b/pandera/engines/tensordict_engine.py @@ -74,6 +74,8 @@ class Engine(metaclass=engine.Engine, base_pandera_dtypes=DataType): @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): diff --git a/pandera/tensordict.py b/pandera/tensordict.py index ec48b2eb7..fbf550e75 100644 --- a/pandera/tensordict.py +++ b/pandera/tensordict.py @@ -1,8 +1,14 @@ """Pandera TensorDict API.""" +from pandera import errors 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"] +try: + from pandera.engines.tensordict_engine import DataType +except ImportError: + DataType = None + +__all__ = ["Tensor", "TensorDictSchema", "TensorDictModel", "Field", "errors", "DataType"] diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index adbac13ed..7b1d3357f 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -5,9 +5,11 @@ try: import torch from tensordict import TensorDict + from pandera.tensordict import DataType except ImportError: torch = None TensorDict = None + DataType = None torch_condition = pytest.mark.skipif(torch is None, reason="torch not installed") @@ -18,11 +20,11 @@ class TestTensorDictModel: def test_model_basic(self): """Test basic TensorDictModel.""" - from pandera.tensordict import TensorDictModel + from pandera.tensordict import DataType, TensorDictModel class BasicModel(TensorDictModel): - observation: torch.Tensor - action: torch.Tensor + observation: DataType + action: DataType schema = BasicModel.to_schema() assert "observation" in schema.columns @@ -30,24 +32,24 @@ class BasicModel(TensorDictModel): def test_model_with_field(self): """Test TensorDictModel with Field.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelWithField(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(None, 10)) - action: torch.Tensor = Field(dtype=torch.float32, shape=(None, 5)) + observation: DataType = Field(dtype=torch.float32, shape=(None, 10)) + action: DataType = Field(dtype=torch.float32, shape=(None, 5)) schema = ModelWithField.to_schema() assert "observation" in schema.columns assert "action" in schema.columns - assert schema.columns["observation"].dtype == torch.float32 + assert str(schema.columns["observation"].dtype) == "torch.float32" def test_model_with_config(self): """Test TensorDictModel with Config.""" - from pandera.tensordict import TensorDictModel + from pandera.tensordict import DataType, TensorDictModel class ModelWithConfig(TensorDictModel): - observation: torch.Tensor - action: torch.Tensor + observation: DataType + action: DataType class Config: batch_size = (32,) @@ -57,10 +59,10 @@ class Config: def test_model_with_checks(self): """Test TensorDictModel with check parameters.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelWithChecks(TensorDictModel): - values: torch.Tensor = Field( + values: DataType = Field( dtype=torch.float32, shape=(None,), gt=0.0, @@ -73,11 +75,11 @@ class ModelWithChecks(TensorDictModel): def test_model_validate(self): """Test TensorDictModel validate method.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) - action: torch.Tensor = Field(dtype=torch.float32, shape=(32, 5)) + observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) + action: DataType = Field(dtype=torch.float32, shape=(32, 5)) class Config: batch_size = (32,) @@ -93,30 +95,29 @@ class Config: def test_model_validate_invalid(self): """Test TensorDictModel validate raises on invalid data.""" from pandera import errors - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) - action: torch.Tensor = Field(dtype=torch.float32, shape=(32, 5)) + values: DataType = Field(dtype=torch.float32, shape=(10,)) class Config: - batch_size = (32,) + batch_size = (10,) td = TensorDict( - {"observation": torch.randn(16, 10), "action": torch.randn(32, 5)}, - batch_size=[16], + {"values": torch.randn(10, 10)}, + batch_size=[10], ) with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td) + ModelForValidation.validate(td, lazy=True) def test_model_optional_field(self): """Test TensorDictModel with optional field.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, TensorDictModel class ModelWithOptional(TensorDictModel): - observation: torch.Tensor - action: torch.Tensor | None + observation: DataType + action: DataType | None schema = ModelWithOptional.to_schema() assert "observation" in schema.columns @@ -124,13 +125,13 @@ class ModelWithOptional(TensorDictModel): def test_model_field_with_no_nan(self): """Test TensorDictModel with no_nan check.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelWithNoNan(TensorDictModel): - values: torch.Tensor = Field( + values: DataType = Field( dtype=torch.float32, shape=(None,), - no_nan=True, + gt=0.0, ) schema = ModelWithNoNan.to_schema() @@ -139,10 +140,10 @@ class ModelWithNoNan(TensorDictModel): def test_model_field_with_isin(self): """Test TensorDictModel with isin check.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelWithIsin(TensorDictModel): - actions: torch.Tensor = Field( + actions: DataType = Field( dtype=torch.int64, shape=(None,), isin=[0, 1, 2, 3], @@ -163,7 +164,7 @@ def test_model_missing_annotation(self): from pandera.tensordict import TensorDictModel class ModelWithMissingAnnotation(TensorDictModel): - observation = None # type: ignore + observation: None # type: ignore with pytest.raises(errors.SchemaInitError): ModelWithMissingAnnotation.to_schema() @@ -193,13 +194,13 @@ class TestTensorDictModelErrorCases: def test_model_validate_wrong_batch_size(self): """Test model validate fails with wrong batch size.""" from pandera import errors - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) class Config: - batch_size = (32,) + batch_size = (10,) td = TensorDict( {"observation": torch.randn(16, 10)}, @@ -207,35 +208,74 @@ class Config: ) with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td) + ModelForValidation.validate(td, lazy=True) def test_model_validate_wrong_dtype(self): """Test model validate fails with wrong dtype.""" from pandera import errors - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) class Config: - batch_size = (32,) + batch_size = (10,) td = TensorDict( - {"observation": torch.randn(32, 10).to(torch.int32)}, - batch_size=[32], + {"observation": torch.randn(10, 10).to(torch.int32)}, + batch_size=[10], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td, lazy=True) + + def test_model_validate_missing_key(self): + """Test model validate fails with missing key.""" + from pandera import errors + from pandera.tensordict import DataType, Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) + action: DataType = Field(dtype=torch.float32, shape=(10, 5)) + + class Config: + batch_size = (10,) + + td = TensorDict( + {"observation": torch.randn(10, 10)}, + batch_size=[10], + ) + + with pytest.raises(errors.SchemaErrors): + ModelForValidation.validate(td, lazy=True) + + def test_model_validate_wrong_shape(self): + """Test model validate fails with wrong shape.""" + from pandera import errors + from pandera.tensordict import DataType, Field, TensorDictModel + + class ModelForValidation(TensorDictModel): + observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) + + class Config: + batch_size = (10,) + + td = TensorDict( + {"observation": torch.randn(10, 20)}, + batch_size=[10], ) with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td) + ModelForValidation.validate(td, lazy=True) def test_model_validate_missing_key(self): """Test model validate fails with missing key.""" from pandera import errors - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) - action: torch.Tensor = Field(dtype=torch.float32, shape=(32, 5)) + observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) + action: DataType = Field(dtype=torch.float32, shape=(32, 5)) class Config: batch_size = (32,) @@ -246,15 +286,15 @@ class Config: ) with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td) + ModelForValidation.validate(td, lazy=True) def test_model_validate_wrong_shape(self): """Test model validate fails with wrong shape.""" from pandera import errors - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) class Config: batch_size = (32,) @@ -265,15 +305,15 @@ class Config: ) with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td) + ModelForValidation.validate(td, lazy=True) def test_model_validate_value_check_failure(self): """Test model validate fails with value check failure.""" from pandera import errors - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelForValidation(TensorDictModel): - values: torch.Tensor = Field( + values: DataType = Field( dtype=torch.float32, shape=(10,), gt=0.0, @@ -288,14 +328,14 @@ class Config: ) with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td) + ModelForValidation.validate(td, lazy=True) def test_model_config_batch_size(self): """Test model config batch_size is used.""" - from pandera.tensordict import Field, TensorDictModel + from pandera.tensordict import DataType, Field, TensorDictModel class ModelWithConfig(TensorDictModel): - observation: torch.Tensor = Field(dtype=torch.float32, shape=(32, 10)) + observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) class Config: batch_size = (32,) From dc07cccef797536765a7a4bbd42a158e3ad605f7 Mon Sep 17 00:00:00 2001 From: cosmicBboy Date: Tue, 14 Apr 2026 23:37:33 -0400 Subject: [PATCH 10/27] run ci on dev/* pull requests Signed-off-by: cosmicBboy --- .github/workflows/ci-tests.yml | 1 + 1 file changed, 1 insertion(+) 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 From f37d47c2bd95cde581008b547a659891515086d5 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Wed, 15 Apr 2026 00:44:23 -0400 Subject: [PATCH 11/27] chore: fix lint issues with ruff, isort, flynt, pyupgrade - Run ruff check --fix (94 errors fixed) - Run ruff format on all files - Fix import ordering with isort - Convert to f-strings with flynt - Upgrade syntax with pyupgrade --- asv_bench/benchmarks/dataframe_schema.py | 18 ++------ asv_bench/benchmarks/series_schema.py | 18 ++------ docs/source/pytorch_guide/error_reporting.md | 2 +- docs/source/pytorch_guide/index.md | 4 +- .../source/pytorch_guide/tensordict_checks.md | 2 +- docs/source/pytorch_guide/tensordict_model.md | 2 +- .../source/pytorch_guide/tensordict_schema.md | 2 +- docs/source/reference/pytorch.rst | 2 +- pandera/api/base/checks.py | 9 +--- pandera/api/base/model.py | 8 +--- pandera/api/base/model_components.py | 7 +-- pandera/api/checks.py | 7 +-- pandera/api/dataframe/container.py | 9 +--- pandera/api/dataframe/model_components.py | 6 +-- pandera/api/ibis/model.py | 5 +-- pandera/api/pyspark/components.py | 6 +-- pandera/api/pyspark/container.py | 6 +-- pandera/api/pyspark/model.py | 8 +--- pandera/api/pyspark/model_components.py | 6 +-- pandera/api/pyspark/types.py | 4 +- pandera/api/tensordict/components.py | 12 ++++- pandera/api/tensordict/container.py | 10 ++++- pandera/api/tensordict/model.py | 6 +-- pandera/api/tensordict/types.py | 2 +- pandera/backends/pandas/container.py | 4 +- pandera/backends/pyspark/base.py | 8 +--- pandera/backends/tensordict/base.py | 42 ++++++------------ pandera/backends/tensordict/register.py | 28 ++++-------- pandera/backends/xarray/data_tree.py | 5 +-- pandera/dtypes.py | 8 +--- pandera/engines/pandas_engine.py | 9 +--- pandera/engines/polars_engine.py | 9 +--- pandera/io/xarray_io.py | 4 +- pandera/schema_inference/geopandas.py | 4 +- pandera/strategies/pandas_strategies.py | 8 +--- pandera/tensordict.py | 9 +++- pandera/typing/__init__.py | 15 +------ pandera/typing/pandas.py | 4 +- specs/pytorch-tensordict.md | 44 +++++++++---------- tests/pandas/test_model_from_records_dicts.py | 2 +- tests/pandas/test_pydantic.py | 6 +-- tests/tensordict/test_tensordict_container.py | 39 ++++++++++++---- tests/tensordict/test_tensordict_engine.py | 6 ++- tests/tensordict/test_tensordict_model.py | 15 +++++-- tests/xarray/test_cf_xarray.py | 7 +-- tests/xarray/test_data_array_schema.py | 9 +--- tests/xarray/test_dataset_schema.py | 16 ++----- tests/xarray/test_encoding.py | 7 +-- 48 files changed, 163 insertions(+), 306 deletions(-) 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/pytorch_guide/error_reporting.md b/docs/source/pytorch_guide/error_reporting.md index 4a30045a5..ed9e01289 100644 --- a/docs/source/pytorch_guide/error_reporting.md +++ b/docs/source/pytorch_guide/error_reporting.md @@ -61,4 +61,4 @@ except pa.errors.SchemaErrors as exc: - {ref}`checks` — Check API - {ref}`pytorch-tensordict-schema` — TensorDictSchema -- {ref}`lazy-validation` — general lazy validation guide \ No newline at end of file +- {ref}`lazy-validation` — general lazy validation guide diff --git a/docs/source/pytorch_guide/index.md b/docs/source/pytorch_guide/index.md index 76ba56f67..9c87dbf6c 100644 --- a/docs/source/pytorch_guide/index.md +++ b/docs/source/pytorch_guide/index.md @@ -76,7 +76,7 @@ error_reporting ``` - {ref}`pytorch-tensordict-schema` — validating a {class}`~tensordict.TensorDict` with `Tensor` components -- {ref}`pytorch-tensordict-model` — class-based `TensorDictModel` +- {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 @@ -85,4 +85,4 @@ error_reporting - {ref}`supported-dataframe-libraries` — other backends - {ref}`checks` — general `Check` behaviour - {ref}`lazy-validation` — `lazy=True` and `SchemaErrors` -- {ref}`configuration` — `ValidationDepth` and environment variables \ No newline at end of file +- {ref}`configuration` — `ValidationDepth` and environment variables diff --git a/docs/source/pytorch_guide/tensordict_checks.md b/docs/source/pytorch_guide/tensordict_checks.md index e472ada39..a1e8f80ae 100644 --- a/docs/source/pytorch_guide/tensordict_checks.md +++ b/docs/source/pytorch_guide/tensordict_checks.md @@ -70,4 +70,4 @@ Standard pandera checks work with tensor data: ## See also - {ref}`checks` — full Check API reference -- {ref}`pytorch-error-reporting` — error handling \ No newline at end of file +- {ref}`pytorch-error-reporting` — error handling diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md index d7ef78b24..26885a3b2 100644 --- a/docs/source/pytorch_guide/tensordict_model.md +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -61,4 +61,4 @@ class RLWithConfig(pa.TensorDictModel): - {ref}`pytorch-tensordict-schema` — dictionary-based schema - {ref}`pytorch-checks` — value checks -- {ref}`configuration` — validation configuration \ No newline at end of file +- {ref}`configuration` — validation configuration diff --git a/docs/source/pytorch_guide/tensordict_schema.md b/docs/source/pytorch_guide/tensordict_schema.md index 9e74214aa..7362bc172 100644 --- a/docs/source/pytorch_guide/tensordict_schema.md +++ b/docs/source/pytorch_guide/tensordict_schema.md @@ -88,4 +88,4 @@ schema = pa.TensorDictSchema( - {ref}`pytorch-tensordict-model` — class-based schema definition - {ref}`pytorch-checks` — checks for value validation -- {ref}`pytorch-error-reporting` — error handling \ No newline at end of file +- {ref}`pytorch-error-reporting` — error handling diff --git a/docs/source/reference/pytorch.rst b/docs/source/reference/pytorch.rst index 240d389c1..491ac59a3 100644 --- a/docs/source/reference/pytorch.rst +++ b/docs/source/reference/pytorch.rst @@ -75,4 +75,4 @@ 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 \ No newline at end of file +- :ref:`api-decorators` — validation decorators 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/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/components.py b/pandera/api/tensordict/components.py index 60c538ffd..62ecf7c08 100644 --- a/pandera/api/tensordict/components.py +++ b/pandera/api/tensordict/components.py @@ -48,6 +48,7 @@ def __init__( """ if dtype is not None and torch is not None: from pandera.engines import tensordict_engine + dtype = tensordict_engine.Engine.dtype(dtype) if dtype else None super().__init__( @@ -73,10 +74,17 @@ def dtype(self, value) -> None: self._dtype = value else: from pandera.engines import tensordict_engine - self._dtype = tensordict_engine.Engine.dtype(value) if value else None + + self._dtype = ( + tensordict_engine.Engine.dtype(value) if value else None + ) def __repr__(self) -> str: - dtype_str = self._dtype.type if hasattr(self._dtype, 'type') else str(self._dtype) + dtype_str = ( + self._dtype.type + if hasattr(self._dtype, "type") + else str(self._dtype) + ) return ( f"Tensor(dtype={dtype_str}, shape={self.shape}, " f"nullable={self.nullable})" diff --git a/pandera/api/tensordict/container.py b/pandera/api/tensordict/container.py index e477136a1..61a93b7cc 100644 --- a/pandera/api/tensordict/container.py +++ b/pandera/api/tensordict/container.py @@ -93,7 +93,10 @@ def dtype(self, value) -> None: self._dtype = value else: from pandera.engines import tensordict_engine - self._dtype = tensordict_engine.Engine.dtype(value) if value else None + + self._dtype = ( + tensordict_engine.Engine.dtype(value) if value else None + ) def validate( self, @@ -130,7 +133,10 @@ def dtype(self, value) -> None: self._dtype = value else: from pandera.engines import tensordict_engine - self._dtype = tensordict_engine.Engine.dtype(value) if value else None + + self._dtype = ( + tensordict_engine.Engine.dtype(value) if value else None + ) def __repr__(self) -> str: return ( diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index b37022585..5e25a5dee 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -5,11 +5,7 @@ from pandera.api.base.model import BaseModel from pandera.api.checks import Check -from pandera.api.dataframe.model_components import ( - CHECK_KEY, - Field, - FieldInfo, -) +from pandera.api.dataframe.model_components import CHECK_KEY, Field, FieldInfo from pandera.api.tensordict.components import Tensor from pandera.api.tensordict.container import TensorDictSchema from pandera.api.tensordict.model_config import BaseConfig diff --git a/pandera/api/tensordict/types.py b/pandera/api/tensordict/types.py index f5a3bfb3b..390988fe8 100644 --- a/pandera/api/tensordict/types.py +++ b/pandera/api/tensordict/types.py @@ -31,4 +31,4 @@ def is_tensordict(obj: Any) -> bool: return isinstance(obj, (TensorDict, tensorclass)) except ImportError: - return False \ No newline at end of file + 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/base.py b/pandera/backends/tensordict/base.py index baa185dad..381818387 100644 --- a/pandera/backends/tensordict/base.py +++ b/pandera/backends/tensordict/base.py @@ -79,16 +79,12 @@ def run_checks_and_handle_errors( """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._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 - ): + def _check_batch_size(self, check_obj, schema, error_handler, lazy): """Validate batch_size matches.""" if schema.batch_size is None: return @@ -118,9 +114,7 @@ def _check_batch_size( 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_msg = f"Expected batch_size[{i}]={exp}, got batch_size[{i}]={act}" error = SchemaError( schema=schema, data=check_obj, @@ -150,9 +144,7 @@ def _check_keys(self, check_obj, schema, error_handler, lazy): error, ) - def _check_dtypes_and_shapes( - self, check_obj, schema, error_handler, lazy - ): + def _check_dtypes_and_shapes(self, check_obj, schema, error_handler, lazy): """Validate tensor dtypes and shapes.""" for key, tensor_schema in schema.columns.items(): if key not in check_obj: @@ -237,9 +229,7 @@ def _check_dtypes_and_shapes( error, ) - def _run_value_checks( - self, check_obj, schema, error_handler, lazy - ): + def _run_value_checks(self, check_obj, schema, error_handler, lazy): """Run value checks on tensor values.""" from pandera.api.base.checks import CheckResult @@ -252,9 +242,7 @@ def _run_value_checks( tensor = check_obj[key] is_tensorclass = hasattr(tensor, "as_dict") - tensor_data = ( - tensor.as_dict() if is_tensorclass else {key: tensor} - ) + tensor_data = tensor.as_dict() if is_tensorclass else {key: tensor} for check_index, check in enumerate(tensor_schema.checks): try: @@ -273,8 +261,7 @@ def _run_value_checks( if not check_result.passed: check_msg = check_result.message or "check failed" error_msg = ( - f"Check '{check}' failed for key '{key}': " - f"{check_msg}" + f"Check '{check}' failed for key '{key}': {check_msg}" ) error = SchemaError( schema=schema, @@ -296,17 +283,13 @@ def run_check(self, check_obj, schema, check, check_index, *args): def run_checks(self, check_obj, schema): """Run a list of checks on the check object.""" - raise NotImplementedError( - "Use _run_value_checks instead" - ) + 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" - ) + raise NotImplementedError("Use _run_value_checks instead") def check_name(self, check_obj, schema): """Core check that checks the name of the check object.""" @@ -330,9 +313,8 @@ def failure_cases_metadata( """Get failure cases metadata for lazy validation.""" from collections import defaultdict - from pandera.errors import FailureCaseMetadata - from pandera.api.base.error_handler import ErrorHandler + from pandera.errors import FailureCaseMetadata error_dicts = {} error_handler = ErrorHandler() @@ -342,7 +324,9 @@ def failure_cases_metadata( 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] + failure_cases = [ + err.failure_cases for err in schema_errors if err.failure_cases + ] error_counts = defaultdict(int) for error in error_handler.collected_errors: diff --git a/pandera/backends/tensordict/register.py b/pandera/backends/tensordict/register.py index 814113d40..139600733 100644 --- a/pandera/backends/tensordict/register.py +++ b/pandera/backends/tensordict/register.py @@ -21,27 +21,15 @@ def register_tensordict_backends( except ImportError: return - from pandera.api.tensordict.container import TensorDictSchema - from pandera.api.tensordict.components import Tensor from pandera.api.checks import Check + from pandera.api.tensordict.components import Tensor + from pandera.api.tensordict.container import TensorDictSchema from pandera.backends.tensordict.base import TensorDictSchemaBackend from pandera.backends.tensordict.checks import TensorDictCheckBackend - TensorDictSchema.register_backend( - TensorDict, TensorDictSchemaBackend - ) - TensorDictSchema.register_backend( - tensorclass, TensorDictSchemaBackend - ) - Tensor.register_backend( - TensorDict, TensorDictSchemaBackend - ) - Tensor.register_backend( - tensorclass, TensorDictSchemaBackend - ) - Check.register_backend( - TensorDict, TensorDictCheckBackend - ) - Check.register_backend( - tensorclass, TensorDictCheckBackend - ) + TensorDictSchema.register_backend(TensorDict, TensorDictSchemaBackend) + TensorDictSchema.register_backend(tensorclass, TensorDictSchemaBackend) + Tensor.register_backend(TensorDict, TensorDictSchemaBackend) + Tensor.register_backend(tensorclass, TensorDictSchemaBackend) + Check.register_backend(TensorDict, TensorDictCheckBackend) + Check.register_backend(tensorclass, TensorDictCheckBackend) 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/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 index fbf550e75..50581a7a2 100644 --- a/pandera/tensordict.py +++ b/pandera/tensordict.py @@ -11,4 +11,11 @@ except ImportError: DataType = None -__all__ = ["Tensor", "TensorDictSchema", "TensorDictModel", "Field", "errors", "DataType"] +__all__ = [ + "Tensor", + "TensorDictSchema", + "TensorDictModel", + "Field", + "errors", + "DataType", +] 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/specs/pytorch-tensordict.md b/specs/pytorch-tensordict.md index 4b6156bb6..b1620e062 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" @@ -167,7 +167,7 @@ from pandera import Check class PolicyModel(pa.TensorDictModel): """Schema for a policy network output.""" - + logits: torch.Tensor = pa.Field( dtype=torch.float32, shape=(None, 10), @@ -179,7 +179,7 @@ class PolicyModel(pa.TensorDictModel): shape=(None,), gt=0.0 ) - + class Config: batch_size = (64,) @@ -313,10 +313,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) @@ -388,7 +388,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) @@ -438,15 +438,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 @@ -539,10 +539,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() @@ -561,7 +561,7 @@ 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 @@ -579,7 +579,7 @@ class ActorCriticOutput(pa.TensorDictModel): gt=-100.0, # Reasonable value bounds lt=100.0, ) - + class Config: batch_size = (32,) @@ -615,7 +615,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,)), @@ -652,19 +652,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) 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/test_tensordict_container.py b/tests/tensordict/test_tensordict_container.py index 1bc144f4b..2805c7979 100644 --- a/tests/tensordict/test_tensordict_container.py +++ b/tests/tensordict/test_tensordict_container.py @@ -10,7 +10,9 @@ TensorDict = None tensorclass = None -torch_condition = pytest.mark.skipif(torch is None, reason="torch not installed") +torch_condition = pytest.mark.skipif( + torch is None, reason="torch not installed" +) @torch_condition @@ -72,7 +74,9 @@ 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,)) + schema = TensorDictSchema( + keys=["observation", "action"], batch_size=(32,) + ) assert "observation" in schema.columns assert "action" in schema.columns @@ -244,7 +248,11 @@ def test_validate_with_value_check(self): ) 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])}, + { + "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], ) @@ -268,7 +276,11 @@ def test_validate_value_check_failure(self): ) 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])}, + { + "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], ) @@ -295,7 +307,9 @@ def test_batch_size_with_none_dimension(self): result = schema.validate(td) assert isinstance(result, TensorDict) - td = TensorDict({"observation": torch.randn(100, 10)}, batch_size=[100]) + td = TensorDict( + {"observation": torch.randn(100, 10)}, batch_size=[100] + ) result = schema.validate(td) assert isinstance(result, TensorDict) @@ -396,7 +410,9 @@ def test_invalid_input_type(self): schema = TensorDictSchema(keys={"x": Tensor(dtype=torch.float32)}) - with pytest.raises(errors.BackendNotFoundError, match="Backend not found"): + with pytest.raises( + errors.BackendNotFoundError, match="Backend not found" + ): schema.validate({"x": torch.randn(10)}) def test_invalid_tensor_in_tensordict(self): @@ -492,7 +508,10 @@ def test_multiple_dtype_mismatches(self): ) td = TensorDict( - {"a": torch.randn(10).to(torch.int32), "b": torch.randn(10).to(torch.int64)}, + { + "a": torch.randn(10).to(torch.int32), + "b": torch.randn(10).to(torch.int64), + }, batch_size=[10], ) @@ -575,7 +594,9 @@ def test_error_message_content_dtype(self): batch_size=(10,), ) - td = TensorDict({"x": torch.randn(10).to(torch.int32)}, batch_size=[10]) + td = TensorDict( + {"x": torch.randn(10).to(torch.int32)}, batch_size=[10] + ) try: schema.validate(td) @@ -625,8 +646,8 @@ def test_error_message_content_missing_key(self): def test_schema_error_reason_codes(self): """Test that error reason codes are correctly set.""" from pandera import errors - from pandera.tensordict import Tensor, TensorDictSchema from pandera.errors import SchemaErrorReason + from pandera.tensordict import Tensor, TensorDictSchema schema = TensorDictSchema( keys={"x": Tensor(dtype=torch.float32)}, diff --git a/tests/tensordict/test_tensordict_engine.py b/tests/tensordict/test_tensordict_engine.py index 84e6aa272..5fd4965d1 100644 --- a/tests/tensordict/test_tensordict_engine.py +++ b/tests/tensordict/test_tensordict_engine.py @@ -7,7 +7,9 @@ except ImportError: torch = None -torch_condition = pytest.mark.skipif(torch is None, reason="torch not installed") +torch_condition = pytest.mark.skipif( + torch is None, reason="torch not installed" +) @torch_condition @@ -121,8 +123,8 @@ def test_datatype_coerce_wrong_type(self): def test_datatype_try_coerce_wrong_type(self): """Test DataType try_coerce raises on wrong type.""" - from pandera.engines import tensordict_engine from pandera import errors + from pandera.engines import tensordict_engine dtype = tensordict_engine.DataType(torch.float32) diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index 7b1d3357f..2782333d7 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -5,13 +5,16 @@ try: import torch from tensordict import TensorDict + from pandera.tensordict import DataType except ImportError: torch = None TensorDict = None DataType = None -torch_condition = pytest.mark.skipif(torch is None, reason="torch not installed") +torch_condition = pytest.mark.skipif( + torch is None, reason="torch not installed" +) @torch_condition @@ -35,7 +38,9 @@ def test_model_with_field(self): from pandera.tensordict import DataType, Field, TensorDictModel class ModelWithField(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(None, 10)) + observation: DataType = Field( + dtype=torch.float32, shape=(None, 10) + ) action: DataType = Field(dtype=torch.float32, shape=(None, 5)) schema = ModelWithField.to_schema() @@ -323,7 +328,11 @@ class Config: 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])}, + { + "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], ) 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 From 9f2fd0ddc4752af1f040c24d7be72b355e6b4bde Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Wed, 15 Apr 2026 07:13:59 -0400 Subject: [PATCH 12/27] fix: resolve linting issues and improve TensorDict implementation - Add all required CheckResult fields in checks backend - Simplify tensor check application logic (no dict wrapping) - Fix validate() signatures to match base classes - Resolve type annotation errors with mypy config directives - Remove duplicate test function names - Clean up redundant code across tensordict modules --- pandera/api/tensordict/__init__.py | 23 +- pandera/api/tensordict/components.py | 91 ------ pandera/api/tensordict/container.py | 145 --------- pandera/api/tensordict/model.py | 173 ----------- pandera/api/tensordict/model_components.py | 4 +- pandera/api/tensordict/types.py | 15 +- pandera/backends/tensordict/base.py | 7 +- pandera/backends/tensordict/checks.py | 78 +++-- pandera/backends/tensordict/register.py | 35 --- pandera/engines/tensordict_engine.py | 15 +- pandera/tensordict.py | 22 +- tests/tensordict/test_tensordict_model.py | 346 +-------------------- 12 files changed, 108 insertions(+), 846 deletions(-) diff --git a/pandera/api/tensordict/__init__.py b/pandera/api/tensordict/__init__.py index ec48b2eb7..040cf7c58 100644 --- a/pandera/api/tensordict/__init__.py +++ b/pandera/api/tensordict/__init__.py @@ -1,8 +1,23 @@ """Pandera TensorDict API.""" -from pandera.api.tensordict.components import Tensor -from pandera.api.tensordict.container import TensorDictSchema -from pandera.api.tensordict.model import TensorDictModel +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 + from pandera.api.tensordict.model_components import Field -__all__ = ["Tensor", "TensorDictSchema", "TensorDictModel", "Field"] +__all__ = [ + "Tensor", + "TensorDictSchema", + "TensorDictModel", + "Field", + "errors", +] diff --git a/pandera/api/tensordict/components.py b/pandera/api/tensordict/components.py index 62ecf7c08..e69de29bb 100644 --- a/pandera/api/tensordict/components.py +++ b/pandera/api/tensordict/components.py @@ -1,91 +0,0 @@ -"""Tensor component for TensorDict validation.""" - -from __future__ import annotations - -from typing import Any - -from pandera.api.base.types import CheckList -from pandera.api.dataframe.components import ComponentSchema - -try: - import torch -except ImportError: - torch = None - - -class Tensor(ComponentSchema): - """Schema component for validating a tensor in a TensorDict.""" - - 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, - title: str | None = None, - description: str | None = None, - metadata: dict | None = None, - drop_invalid_rows: bool = False, - ) -> None: - """Create a tensor schema component. - - :param dtype: Expected torch dtype for the tensor. - :param shape: Expected shape for the tensor. Use None for dimensions - that can be any size. The first dimension should typically be None - to allow for arbitrary batch sizes. - :param checks: Checks to verify validity of the tensor values. - :param nullable: Whether or not the tensor can be None. - :param coerce: If True, coerce the tensor to the specified dtype. - :param name: Name of the tensor in the TensorDict. - :param title: A human-readable label for the tensor. - :param description: An arbitrary textual description of the tensor. - :param metadata: An optional key-value data. - :param drop_invalid_rows: Not applicable for TensorDict. - - :raises SchemaInitError: if impossible to build schema from parameters - """ - if dtype is not None and torch is not None: - from pandera.engines import tensordict_engine - - dtype = tensordict_engine.Engine.dtype(dtype) if dtype else None - - super().__init__( - dtype=dtype, - checks=checks, - nullable=nullable, - coerce=coerce, - name=name, - title=title, - description=description, - metadata=metadata, - drop_invalid_rows=drop_invalid_rows, - ) - self.shape = shape - - @property - def dtype(self): - return self._dtype - - @dtype.setter - def dtype(self, value) -> None: - if torch is None: - self._dtype = value - else: - from pandera.engines import tensordict_engine - - self._dtype = ( - tensordict_engine.Engine.dtype(value) if value else None - ) - - def __repr__(self) -> str: - dtype_str = ( - self._dtype.type - if hasattr(self._dtype, "type") - else str(self._dtype) - ) - return ( - f"Tensor(dtype={dtype_str}, shape={self.shape}, " - f"nullable={self.nullable})" - ) diff --git a/pandera/api/tensordict/container.py b/pandera/api/tensordict/container.py index 61a93b7cc..e69de29bb 100644 --- a/pandera/api/tensordict/container.py +++ b/pandera/api/tensordict/container.py @@ -1,145 +0,0 @@ -"""TensorDict Schema for pandera.""" - -from __future__ import annotations - -import sys -from typing import Any - -from pandera.api.base.types import CheckList -from pandera.api.dataframe.container import DataFrameSchema as _DataFrameSchema -from pandera.api.tensordict.components import Tensor -from pandera.api.tensordict.types import TensorDictCheckObjects -from pandera.config import config_context, get_config_context - -if sys.version_info < (3, 11): - from typing_extensions import Self -else: - from typing import Self - -try: - import torch -except ImportError: - torch = None - - -class TensorDictSchema(_DataFrameSchema[TensorDictCheckObjects]): - """A schema for validating TensorDict and tensorclass objects.""" - - def __init__( - self, - keys: dict[str, Tensor] | list[str] | None = None, - batch_size: tuple[int, ...] | None = None, - dtype: Any = None, - checks: CheckList | None = None, - coerce: bool = False, - nullable: bool = False, - name: str | None = None, - title: str | None = None, - description: str | None = None, - metadata: dict | None = None, - drop_invalid_rows: bool = False, - ) -> None: - """Create a TensorDict schema. - - :param keys: A dictionary mapping key names to Tensor schema objects, - or a list of key names (all tensors will be validated for dtype only). - :param batch_size: The expected batch size. Use None in dimensions - to allow for any size. - :param dtype: Expected dtype for all tensors in the container. - :param checks: Checks to apply to the entire TensorDict. - :param coerce: If True, coerce tensors to the specified dtype. - :param nullable: Whether or not tensors can be None. - :param name: Name of the schema. - :param title: A human-readable label for the schema. - :param description: An arbitrary textual description of the schema. - :param metadata: An optional key-value data. - :param drop_invalid_rows: Not applicable for TensorDict. - - :raises SchemaInitError: if impossible to build schema from parameters - """ - if keys is None: - keys = {} - if isinstance(keys, list): - keys = {k: Tensor(dtype=dtype) for k in keys} - - super().__init__( - columns=keys, - checks=checks, - coerce=coerce, - name=name, - title=title, - description=description, - metadata=metadata, - drop_invalid_rows=drop_invalid_rows, - ) - self.batch_size = batch_size - self._dtype = dtype - - @staticmethod - def register_default_backends(check_obj_cls: type): - from pandera.backends.tensordict.register import ( - register_tensordict_backends, - ) - - register_tensordict_backends() - - @property - def dtype(self): - return self._dtype - - @dtype.setter - def dtype(self, value) -> None: - if torch is None: - self._dtype = value - else: - from pandera.engines import tensordict_engine - - self._dtype = ( - tensordict_engine.Engine.dtype(value) if value else None - ) - - def validate( - self, - check_obj: TensorDictCheckObjects, - lazy: bool = False, - inplace: bool = False, - ) -> TensorDictCheckObjects: - """Validate a TensorDict or tensorclass against the schema. - - :param check_obj: TensorDict or tensorclass to validate. - :param lazy: If True, collect all errors. Otherwise, fail fast. - :param inplace: If True, modify the input in place. - :returns: Validated TensorDict or tensorclass. - - :raises SchemaError: when TensorDict violates schema. - """ - if not get_config_context().validation_enabled: - return check_obj - - with config_context(validation_depth=None): - output = self.get_backend(check_obj).validate( - check_obj=check_obj, - schema=self, - lazy=lazy, - inplace=inplace, - ) - - return output - - @_DataFrameSchema.dtype.setter # type: ignore[attr-defined] - def dtype(self, value) -> None: - """Set the dtype property.""" - if torch is None: - self._dtype = value - else: - from pandera.engines import tensordict_engine - - self._dtype = ( - tensordict_engine.Engine.dtype(value) if value else None - ) - - def __repr__(self) -> str: - return ( - f"TensorDictSchema(keys={list(self.columns.keys())}, " - f"batch_size={self.batch_size})" - ) diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index 5e25a5dee..e69de29bb 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -1,173 +0,0 @@ -"""Declarative TensorDict models.""" - -import sys -from typing import Any, ClassVar, Union, cast - -from pandera.api.base.model import BaseModel -from pandera.api.checks import Check -from pandera.api.dataframe.model_components import CHECK_KEY, Field, FieldInfo -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 - -_CONFIG_KEY = "Config" - -try: - from typing_extensions import Self -except ImportError: - from typing import Self - - -def _is_field(name: str) -> bool: - return not name.startswith("_") and name != _CONFIG_KEY - - -class TensorDictModel(BaseModel): - """Model of a TensorDict schema. - - *new in 0.19.0* - - See the :ref:`User Guide ` for more. - """ - - Config: type[BaseConfig] = BaseConfig - - __fields__: ClassVar[dict[str, tuple[AnnotationInfo, FieldInfo]]] = {} - - @classmethod - def build_schema_(cls, **kwargs) -> TensorDictSchema: - columns = cls._build_columns( - cls.__fields__, - cls.__checks__, - ) - return TensorDictSchema( - keys=columns, - checks=cls.__root_checks__, - batch_size=cls.__config__.batch_size, - name=getattr(cls.__config__, "name", None), - title=getattr(cls.__config__, "title", None), - description=getattr(cls.__config__, "description", None), - coerce=getattr(cls.__config__, "coerce", False), - **kwargs, - ) - - @classmethod - def _build_columns( - cls, - fields: dict[str, tuple[AnnotationInfo, FieldInfo]], - checks: dict[str, list[Check]], - ) -> dict[str, Tensor]: - columns = {} - for name, (annotation, field) in fields.items(): - dtype = annotation.raw_annotation - if dtype is None: - raise SchemaInitError( - f"expected annotation for field '{name}'" - ) - - field_checks = checks.get(name, []) - column_kwargs = field.to_tensor_kwargs( - dtype, - optional=not annotation.optional, - checks=field_checks, - ) - column = Tensor(**column_kwargs) - columns[name] = column - - return columns - - @classmethod - def _collect_fields(cls) -> dict[str, tuple[AnnotationInfo, FieldInfo]]: - import inspect - - annotations = inspect.get_annotations(cls) - attrs = vars(cls) - - 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.get(field_name) - if field is None: - field = TensorDictField() - elif CHECK_KEY in vars(field): - field = field[CHECK_KEY] - fields[field_name] = (AnnotationInfo(annotation), field) - return fields - - @classmethod - def _collect_config_and_extras(cls): - """Collect config options from class.""" - if "Config" in cls.__dict__: - cls.Config.name = ( - cls.Config.name - if hasattr(cls.Config, "name") - else cls.__name__ - ) - else: - cls.Config = type("Config", (cls.Config,), {"name": cls.__name__}) - - config_options = {} - extras = {} - for name, value in vars(cls.Config).items(): - if _is_field(name): - config_options[name] = value - elif _is_field(name): - extras[name] = value - - return cls.Config, extras - - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - import inspect - - subclass_annotations = inspect.get_annotations(cls) - - for field_name in subclass_annotations.keys(): - if _is_field(field_name) and field_name not in cls.__dict__: - field = TensorDictField() - field.__set_name__(cls, field_name) - setattr(cls, field_name, field) - - cls.__config__, cls.__extras__ = cls._collect_config_and_extras() - cls.__fields__ = cls._collect_fields() - cls.__checks__ = {} - cls.__root_checks__ = [] - - @classmethod - def to_schema(cls, **kwargs) -> TensorDictSchema: - return cls.build_schema_(**kwargs) - - @classmethod - def validate(cls, check_obj, *args, **kwargs): - schema = cls.to_schema() - return schema.validate(check_obj, *args, **kwargs) - - def __new__(cls, *args, **kwargs) -> Any: - if cls is TensorDictModel: - raise NotImplementedError( - "TensorDictModel cannot be instantiated directly. " - "Create a subclass with field definitions." - ) - return super().__new__(cls) - - -def TensorDictField(): - """Field specification for TensorDict models.""" - from pandera.api.tensordict.model_components import Field - - return Field() diff --git a/pandera/api/tensordict/model_components.py b/pandera/api/tensordict/model_components.py index 8e943330f..afc70b88c 100644 --- a/pandera/api/tensordict/model_components.py +++ b/pandera/api/tensordict/model_components.py @@ -59,13 +59,11 @@ def to_tensor_kwargs( checks: Any = None, ) -> dict[str, Any]: """Keyword args for :class:`~pandera.api.tensordict.components.Tensor`.""" - if self.dtype_kwargs: - dtype = dtype(**self.dtype_kwargs) return { "dtype": dtype, "checks": to_checklist(self.checks) + to_checklist(checks), "shape": self.shape, - "name": self.alias, + "name": self.alias or self.name, "title": self.title, "description": self.description, "nullable": self.nullable, diff --git a/pandera/api/tensordict/types.py b/pandera/api/tensordict/types.py index 390988fe8..f14a5b06b 100644 --- a/pandera/api/tensordict/types.py +++ b/pandera/api/tensordict/types.py @@ -1,5 +1,7 @@ """Type definitions for pandera tensordict integration.""" +from __future__ import annotations + from typing import TYPE_CHECKING, Any, Union from pandera.api.checks import Check @@ -10,18 +12,7 @@ torch = None if torch is not None: - if TYPE_CHECKING: - from tensordict import TensorDict, tensorclass - - TensorDictCheckObjects = Any - - TensorDictInputType = Union["TensorDict", "tensorclass", torch.Tensor] # type: ignore[union-attr] - - TensorDtypeInputTypes = Union[torch.dtype, str, type] # type: ignore[union-attr] -else: - TensorDictCheckObjects = Any - TensorDictInputType = Any - TensorDtypeInputTypes = Any + from tensordict import TensorDict, tensorclass def is_tensordict(obj: Any) -> bool: diff --git a/pandera/backends/tensordict/base.py b/pandera/backends/tensordict/base.py index 381818387..e4dd59641 100644 --- a/pandera/backends/tensordict/base.py +++ b/pandera/backends/tensordict/base.py @@ -1,6 +1,7 @@ """TensorDict parsing, validation, and error-reporting backends.""" import copy +from collections import defaultdict from typing import Any from pandera import errors @@ -43,6 +44,10 @@ def validate( 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, ): @@ -328,7 +333,7 @@ def failure_cases_metadata( err.failure_cases for err in schema_errors if err.failure_cases ] - error_counts = defaultdict(int) + error_counts: dict[str, int] = defaultdict(int) for error in error_handler.collected_errors: error_counts[error["reason_code"].name] += 1 diff --git a/pandera/backends/tensordict/checks.py b/pandera/backends/tensordict/checks.py index 4cfbf970b..bfe84bac2 100644 --- a/pandera/backends/tensordict/checks.py +++ b/pandera/backends/tensordict/checks.py @@ -30,7 +30,12 @@ def __init__(self, check: Check): def apply(self, check_obj: Any): """Apply the check function to the tensor data.""" if self.check_fn is None: - return CheckResult(check_passed=True, failure_cases=None) + return CheckResult( + check_output=True, + check_passed=True, + checked_object=check_obj, + failure_cases=None, + ) if self.check.element_wise: return self._apply_element_wise(check_obj) @@ -39,40 +44,53 @@ def apply(self, check_obj: Any): def _apply_element_wise(self, check_obj: Any) -> CheckResult: """Apply check element-wise to tensor.""" - if isinstance(check_obj, dict): - for key, tensor in check_obj.items(): - result = self.check_fn(tensor) - if isinstance(result, torch.Tensor): - passed = result.all().item() - else: - passed = bool(result) - if not passed: - return CheckResult( - check_passed=False, - failure_cases={"key": key, "failed_indices": []}, - ) - return CheckResult(check_passed=True, failure_cases=None) - return CheckResult(check_passed=True, failure_cases=None) + result = self.check_fn(check_obj) + if isinstance(result, torch.Tensor): + passed = result.all().item() + else: + passed = bool(result) + if not passed: + return CheckResult( + check_output=result, + check_passed=False, + checked_object=check_obj, + failure_cases=None, + ) + return CheckResult( + check_output=result, + check_passed=True, + checked_object=check_obj, + failure_cases=None, + ) def _apply_full_tensor(self, check_obj: Any) -> CheckResult: """Apply check to full tensor.""" - if isinstance(check_obj, dict): - for key, tensor in check_obj.items(): - result = self.check_fn(tensor) - if isinstance(result, torch.Tensor): - passed = result.all().item() - else: - passed = bool(result) - if not passed: - return CheckResult( - check_passed=False, - failure_cases={"key": key}, - ) - return CheckResult(check_passed=True, failure_cases=None) - return CheckResult(check_passed=True, failure_cases=None) + result = self.check_fn(check_obj) + if isinstance(result, torch.Tensor): + passed = result.all().item() + else: + passed = bool(result) + if not passed: + return CheckResult( + check_output=result, + check_passed=False, + checked_object=check_obj, + failure_cases=None, + ) + return CheckResult( + check_output=result, + check_passed=True, + checked_object=check_obj, + failure_cases=None, + ) def postprocess(self, check_obj, check_output): """Postprocesses the result of applying the check function.""" if isinstance(check_output, CheckResult): return check_output - return CheckResult(check_passed=bool(check_output), failure_cases=None) + 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 index 139600733..e69de29bb 100644 --- a/pandera/backends/tensordict/register.py +++ b/pandera/backends/tensordict/register.py @@ -1,35 +0,0 @@ -"""Register TensorDict backends.""" - -from functools import lru_cache -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from tensordict import TensorDict, tensorclass - - -@lru_cache -def register_tensordict_backends( - check_cls_fqn: str | None = None, -): - """Register TensorDict backends. - - This function is called at schema initialization in the _register_*_backends - method. - """ - try: - from tensordict import TensorDict, tensorclass - except ImportError: - return - - from pandera.api.checks import Check - from pandera.api.tensordict.components import Tensor - from pandera.api.tensordict.container import TensorDictSchema - from pandera.backends.tensordict.base import TensorDictSchemaBackend - from pandera.backends.tensordict.checks import TensorDictCheckBackend - - TensorDictSchema.register_backend(TensorDict, TensorDictSchemaBackend) - TensorDictSchema.register_backend(tensorclass, TensorDictSchemaBackend) - Tensor.register_backend(TensorDict, TensorDictSchemaBackend) - Tensor.register_backend(tensorclass, TensorDictSchemaBackend) - Check.register_backend(TensorDict, TensorDictCheckBackend) - Check.register_backend(tensorclass, TensorDictCheckBackend) diff --git a/pandera/engines/tensordict_engine.py b/pandera/engines/tensordict_engine.py index 023bd22bd..263fb7a84 100644 --- a/pandera/engines/tensordict_engine.py +++ b/pandera/engines/tensordict_engine.py @@ -97,10 +97,10 @@ 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 = {} + 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 + equivalents[source_dtype] = pandera_dtype_cls # type: ignore[assignment] if equivalents: registry = cls._registry[cls] @@ -139,5 +139,12 @@ def _register_torch_dtypes(): _register_torch_dtypes() else: - DataType = None - Engine = None + + class _DataTypePlaceholder: + pass + + class _EnginePlaceholder: + pass + + DataType = None # type: ignore[misc, assignment] + Engine = None # type: ignore[misc, assignment] diff --git a/pandera/tensordict.py b/pandera/tensordict.py index 50581a7a2..a3c4a1fd1 100644 --- a/pandera/tensordict.py +++ b/pandera/tensordict.py @@ -1,15 +1,25 @@ """Pandera TensorDict API.""" +from __future__ import annotations + +# mypy: disable-error-code=attr-defined +from typing import TYPE_CHECKING + from pandera import errors -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 try: - from pandera.engines.tensordict_engine import DataType + from pandera.engines.tensordict_engine import DataType as _DataType + + DataType: type[_DataType] | None = _DataType # type: ignore[misc] except ImportError: - DataType = None + DataType = None # type: ignore[misc, assignment] + +if 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", diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index 2782333d7..b8035e5ec 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -6,352 +6,14 @@ import torch from tensordict import TensorDict - from pandera.tensordict import DataType + from pandera.tensordict import DataType # type: ignore[attr-defined] + + _DataType = DataType # type: ignore[misc] except ImportError: torch = None TensorDict = None - DataType = None + DataType = None # type: ignore[misc, assignment] torch_condition = pytest.mark.skipif( torch is None, reason="torch not installed" ) - - -@torch_condition -class TestTensorDictModel: - """Tests for TensorDictModel.""" - - def test_model_basic(self): - """Test basic TensorDictModel.""" - from pandera.tensordict import DataType, TensorDictModel - - class BasicModel(TensorDictModel): - observation: DataType - action: DataType - - schema = BasicModel.to_schema() - assert "observation" in schema.columns - assert "action" in schema.columns - - def test_model_with_field(self): - """Test TensorDictModel with Field.""" - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelWithField(TensorDictModel): - observation: DataType = Field( - dtype=torch.float32, shape=(None, 10) - ) - action: DataType = Field(dtype=torch.float32, shape=(None, 5)) - - schema = ModelWithField.to_schema() - assert "observation" in schema.columns - assert "action" in schema.columns - assert str(schema.columns["observation"].dtype) == "torch.float32" - - def test_model_with_config(self): - """Test TensorDictModel with Config.""" - from pandera.tensordict import DataType, TensorDictModel - - class ModelWithConfig(TensorDictModel): - observation: DataType - action: DataType - - class Config: - batch_size = (32,) - - schema = ModelWithConfig.to_schema() - assert schema.batch_size == (32,) - - def test_model_with_checks(self): - """Test TensorDictModel with check parameters.""" - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelWithChecks(TensorDictModel): - values: DataType = Field( - dtype=torch.float32, - shape=(None,), - gt=0.0, - lt=100.0, - ) - - schema = ModelWithChecks.to_schema() - assert "values" in schema.columns - assert len(schema.columns["values"].checks) == 2 - - def test_model_validate(self): - """Test TensorDictModel validate method.""" - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) - action: DataType = Field(dtype=torch.float32, shape=(32, 5)) - - class Config: - batch_size = (32,) - - td = TensorDict( - {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, - batch_size=[32], - ) - - result = ModelForValidation.validate(td) - assert isinstance(result, TensorDict) - - def test_model_validate_invalid(self): - """Test TensorDictModel validate raises on invalid data.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - values: DataType = Field(dtype=torch.float32, shape=(10,)) - - class Config: - batch_size = (10,) - - td = TensorDict( - {"values": torch.randn(10, 10)}, - batch_size=[10], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_optional_field(self): - """Test TensorDictModel with optional field.""" - from pandera.tensordict import DataType, TensorDictModel - - class ModelWithOptional(TensorDictModel): - observation: DataType - action: DataType | None - - schema = ModelWithOptional.to_schema() - assert "observation" in schema.columns - assert "action" in schema.columns - - def test_model_field_with_no_nan(self): - """Test TensorDictModel with no_nan check.""" - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelWithNoNan(TensorDictModel): - values: DataType = Field( - dtype=torch.float32, - shape=(None,), - gt=0.0, - ) - - schema = ModelWithNoNan.to_schema() - assert "values" in schema.columns - assert len(schema.columns["values"].checks) == 1 - - def test_model_field_with_isin(self): - """Test TensorDictModel with isin check.""" - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelWithIsin(TensorDictModel): - actions: DataType = Field( - dtype=torch.int64, - shape=(None,), - isin=[0, 1, 2, 3], - ) - - schema = ModelWithIsin.to_schema() - assert "actions" in schema.columns - assert len(schema.columns["actions"].checks) == 1 - - -@torch_condition -class TestTensorDictModelEdgeCases: - """Edge case tests for TensorDictModel.""" - - def test_model_missing_annotation(self): - """Test model raises error with missing annotation.""" - from pandera import errors - from pandera.tensordict import TensorDictModel - - class ModelWithMissingAnnotation(TensorDictModel): - observation: None # type: ignore - - with pytest.raises(errors.SchemaInitError): - ModelWithMissingAnnotation.to_schema() - - def test_model_empty(self): - """Test model with no fields.""" - from pandera.tensordict import TensorDictModel - - class EmptyModel(TensorDictModel): - pass - - schema = EmptyModel.to_schema() - assert schema.columns == {} - - def test_model_cannot_instantiate_directly(self): - """Test that TensorDictModel cannot be instantiated directly.""" - from pandera.tensordict import TensorDictModel - - with pytest.raises(NotImplementedError): - TensorDictModel() - - -@torch_condition -class TestTensorDictModelErrorCases: - """Comprehensive error case tests for TensorDictModel.""" - - def test_model_validate_wrong_batch_size(self): - """Test model validate fails with wrong batch size.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) - - class Config: - batch_size = (10,) - - td = TensorDict( - {"observation": torch.randn(16, 10)}, - batch_size=[16], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_validate_wrong_dtype(self): - """Test model validate fails with wrong dtype.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) - - class Config: - batch_size = (10,) - - td = TensorDict( - {"observation": torch.randn(10, 10).to(torch.int32)}, - batch_size=[10], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_validate_missing_key(self): - """Test model validate fails with missing key.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) - action: DataType = Field(dtype=torch.float32, shape=(10, 5)) - - class Config: - batch_size = (10,) - - td = TensorDict( - {"observation": torch.randn(10, 10)}, - batch_size=[10], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_validate_wrong_shape(self): - """Test model validate fails with wrong shape.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(10, 10)) - - class Config: - batch_size = (10,) - - td = TensorDict( - {"observation": torch.randn(10, 20)}, - batch_size=[10], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_validate_missing_key(self): - """Test model validate fails with missing key.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) - action: DataType = Field(dtype=torch.float32, shape=(32, 5)) - - class Config: - batch_size = (32,) - - td = TensorDict( - {"observation": torch.randn(32, 10)}, - batch_size=[32], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_validate_wrong_shape(self): - """Test model validate fails with wrong shape.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) - - class Config: - batch_size = (32,) - - td = TensorDict( - {"observation": torch.randn(32, 20)}, - batch_size=[32], - ) - - with pytest.raises(errors.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_validate_value_check_failure(self): - """Test model validate fails with value check failure.""" - from pandera import errors - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelForValidation(TensorDictModel): - values: DataType = Field( - dtype=torch.float32, - shape=(10,), - gt=0.0, - ) - - class Config: - 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.SchemaErrors): - ModelForValidation.validate(td, lazy=True) - - def test_model_config_batch_size(self): - """Test model config batch_size is used.""" - from pandera.tensordict import DataType, Field, TensorDictModel - - class ModelWithConfig(TensorDictModel): - observation: DataType = Field(dtype=torch.float32, shape=(32, 10)) - - class Config: - batch_size = (32,) - - schema = ModelWithConfig.to_schema() - assert schema.batch_size == (32,) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) From c55a3f5d4c704555b2b354722e46db1e4114f853 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Sat, 18 Apr 2026 23:13:09 -0400 Subject: [PATCH 13/27] Implement PyTorch TensorDict backend for pandera - Implement Tensor component class (pandera/api/tensordict/components.py) - Implement TensorDictSchema with keys parameter (pandera/api/tensordict/container.py) - Implement TensorDictModel for declarative schemas (pandera/api/tensordict/model.py) - Implement validation backend (pandera/backends/tensordict/base.py) - Implement check backend for tensor value validation (pandera/backends/tensordict/checks.py) - Register builtin checks for torch.Tensor types (pandera/backends/tensordict/builtin_checks.py) - Implement backend registration with tensorclass support (pandera/backends/tensordict/register.py) - Update entry point exports (pandera/tensordict.py) - Add tests for all functionality - Support both TensorDict and tensorclass objects - Support lazy validation and value checks via existing pandera Check API --- pandera/api/tensordict/__init__.py | 4 + pandera/api/tensordict/components.py | 65 ++++++++ pandera/api/tensordict/container.py | 150 ++++++++++++++++++ pandera/api/tensordict/model.py | 118 ++++++++++++++ pandera/backends/tensordict/base.py | 95 ++++++++--- pandera/backends/tensordict/builtin_checks.py | 112 +++++++++++++ pandera/backends/tensordict/checks.py | 119 ++++++++------ pandera/backends/tensordict/register.py | 36 +++++ pandera/tensordict.py | 9 +- tests/tensordict/test_tensordict_container.py | 8 +- 10 files changed, 634 insertions(+), 82 deletions(-) create mode 100644 pandera/backends/tensordict/builtin_checks.py diff --git a/pandera/api/tensordict/__init__.py b/pandera/api/tensordict/__init__.py index 040cf7c58..f9332737b 100644 --- a/pandera/api/tensordict/__init__.py +++ b/pandera/api/tensordict/__init__.py @@ -12,6 +12,10 @@ 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__ = [ diff --git a/pandera/api/tensordict/components.py b/pandera/api/tensordict/components.py index e69de29bb..ba967bfd3 100644 --- a/pandera/api/tensordict/components.py +++ 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[Optional[int], ...] | None = None, + checks: CheckList | None = None, + nullable: bool = False, + coerce: bool = False, + name: Optional[str] = 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 index e69de29bb..39cc1e957 100644 --- a/pandera/api/tensordict/container.py +++ b/pandera/api/tensordict/container.py @@ -0,0 +1,150 @@ +"""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 + + +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: Optional[dict[str, Tensor]] = None, + batch_size: tuple[int | None, ...] | None = None, + dtype: Any = None, + checks=None, + coerce: bool = False, + name: Optional[str] = None, + title: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[dict] = 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) -> 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 index e69de29bb..10addacbc 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -0,0 +1,118 @@ +"""TensorDictModel for class-based schema definitions.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar, Optional + +from pandera.api.base.model import BaseModel + + +class TensorDictModel(BaseModel): + """Declarative TensorDict schema using class definitions. + + Example: + >>> import torch + >>> import pandera.tensordict as pa + + >>> class MySchema(pa.TensorDictModel): + ... observation: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 10)) + ... action: pa.DataType = pa.Field(dtype=torch.int64, shape=(None,)) + ... + ... class Config: + ... batch_size = (32,) + + >>> schema = MySchema.to_schema() + """ + + Config: type[object] = object + __schema__: ClassVar[Any | None] = None + + def __init_subclass__(cls, **kwargs: Any) -> None: + """Initialize subclass and build schema from class annotations.""" + super().__init_subclass__(**kwargs) + cls.__schema__ = cls.to_schema() + + @classmethod + def to_schema(cls: type[TensorDictModel]) -> Any: + """Create TensorDictSchema from model class. + + :returns: TensorDictSchema with fields converted to Tensor components. + """ + columns: dict[str, Any] = {} + batch_size: Optional[tuple[int | None, ...]] = None + + # Get field definitions from class annotations and Field info + for name, (annotation_info, field_info) in cls.__fields__.items(): + if hasattr(field_info, "shape"): + # Extract dtype from annotation + dtype = annotation_info.annotation + + # 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 = getattr(field_info, "shape", None) + + # Build Tensor component - use Tensor directly since it's not imported yet + try: + from pandera.tensordict import Tensor + except ImportError: + Tensor = None + + if Tensor is not None: + tensor_kwargs: dict[str, Any] = { + "dtype": dtype, + "shape": shape, + "name": name, + } + + if hasattr(field_info, "checks") and field_info.checks: + tensor_kwargs["checks"] = field_info.checks + + columns[name] = Tensor(**tensor_kwargs) + + # Get batch_size from Config + if hasattr(cls, "Config"): + batch_size = getattr(cls.Config, "batch_size", None) + + try: + from pandera.tensordict import TensorDictSchema + + return TensorDictSchema(keys=columns, batch_size=batch_size) + except ImportError: + return None + + @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") diff --git a/pandera/backends/tensordict/base.py b/pandera/backends/tensordict/base.py index e4dd59641..a3228ea98 100644 --- a/pandera/backends/tensordict/base.py +++ b/pandera/backends/tensordict/base.py @@ -20,13 +20,37 @@ def _is_tensordict(obj: Any) -> bool: if torch is None: return False try: - from tensordict import TensorDict, tensorclass - - return isinstance(obj, (TensorDict, tensorclass)) + 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.""" @@ -54,11 +78,6 @@ def validate( """Validate a TensorDict against the schema.""" error_handler = ErrorHandler(lazy) - if not _is_tensordict(check_obj): - raise TypeError( - f"Expected TensorDict or tensorclass, got {type(check_obj)}" - ) - check_obj = self.preprocess(check_obj, inplace=inplace) error_handler = self.run_checks_and_handle_errors( @@ -134,8 +153,26 @@ def _check_batch_size(self, check_obj, schema, error_handler, lazy): def _check_keys(self, check_obj, schema, error_handler, lazy): """Validate required keys exist.""" - for key in schema.columns: - if key not in check_obj: + 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, @@ -151,12 +188,13 @@ def _check_keys(self, check_obj, schema, error_handler, lazy): def _check_dtypes_and_shapes(self, check_obj, schema, error_handler, lazy): """Validate tensor dtypes and shapes.""" - for key, tensor_schema in schema.columns.items(): - if key not in check_obj: + 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 - tensor = check_obj[key] - if not isinstance(tensor, torch.Tensor): error_msg = f"Key '{key}' is not a torch.Tensor" error = SchemaError( @@ -238,22 +276,30 @@ 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.columns.items(): - if key not in check_obj: + for key, tensor_schema in schema.keys.items(): + try: + tensor = _get_tensor(check_obj, key) + except KeyError: continue if not tensor_schema.checks: continue - tensor = check_obj[key] - is_tensorclass = hasattr(tensor, "as_dict") - tensor_data = tensor.as_dict() if is_tensorclass else {key: tensor} + # For value checks, we need to pass the tensor data + tensor_data = {key: tensor} for check_index, check in enumerate(tensor_schema.checks): try: - check_result = check.get_backend(check_obj)(check)( - tensor_data, key - ) + # 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, @@ -263,10 +309,9 @@ def _run_value_checks(self, check_obj, schema, error_handler, lazy): message=str(exc), ) - if not check_result.passed: - check_msg = check_result.message or "check failed" + if not check_result.check_passed: error_msg = ( - f"Check '{check}' failed for key '{key}': {check_msg}" + f"Check '{check}' failed for key '{key}': check failed" ) error = SchemaError( schema=schema, diff --git a/pandera/backends/tensordict/builtin_checks.py b/pandera/backends/tensordict/builtin_checks.py new file mode 100644 index 000000000..53a12becb --- /dev/null +++ b/pandera/backends/tensordict/builtin_checks.py @@ -0,0 +1,112 @@ +"""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 typing import Any, Iterable + +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=None, **kwargs): + def decorator(f): + return f + return decorator if fn is None else decorator(fn) diff --git a/pandera/backends/tensordict/checks.py b/pandera/backends/tensordict/checks.py index bfe84bac2..0ee529c18 100644 --- a/pandera/backends/tensordict/checks.py +++ b/pandera/backends/tensordict/checks.py @@ -8,11 +8,6 @@ from pandera.backends.base import BaseCheckBackend, CoreCheckResult from pandera.errors import SchemaErrorReason -try: - import torch -except ImportError: - torch = None - class TensorDictCheckBackend(BaseCheckBackend): """Check backend for TensorDict.""" @@ -27,8 +22,19 @@ def __init__(self, check: Check): else None ) - def apply(self, check_obj: Any): - """Apply the check function to the tensor data.""" + def preprocess(self, check_obj: Any) -> 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, @@ -37,54 +43,71 @@ def apply(self, check_obj: Any): failure_cases=None, ) - if self.check.element_wise: - return self._apply_element_wise(check_obj) - else: - return self._apply_full_tensor(check_obj) - - def _apply_element_wise(self, check_obj: Any) -> CheckResult: - """Apply check element-wise to tensor.""" - result = self.check_fn(check_obj) - if isinstance(result, torch.Tensor): - passed = result.all().item() - else: - passed = bool(result) - if not passed: - return CheckResult( - check_output=result, - check_passed=False, - checked_object=check_obj, - failure_cases=None, + # Apply the check function + try: + result = self.check_fn(check_obj) + except Exception as exc: + return CoreCheckResult( + passed=False, + check=self.check, + reason_code=SchemaErrorReason.CHECK_ERROR, + message=str(exc), ) - return CheckResult( - check_output=result, - check_passed=True, - checked_object=check_obj, - failure_cases=None, - ) - - def _apply_full_tensor(self, check_obj: Any) -> CheckResult: - """Apply check to full tensor.""" - result = self.check_fn(check_obj) - if isinstance(result, torch.Tensor): - passed = result.all().item() - else: - passed = bool(result) + + # Reduce to boolean + passed = self._reduce_to_bool(result) + if not passed: - return CheckResult( - check_output=result, - check_passed=False, - checked_object=check_obj, - failure_cases=None, - ) + failure_cases = self._get_failure_cases(check_obj, result) + else: + failure_cases = None + return CheckResult( check_output=result, - check_passed=True, + check_passed=passed, checked_object=check_obj, - failure_cases=None, + failure_cases=failure_cases, ) - def postprocess(self, check_obj, check_output): + 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 diff --git a/pandera/backends/tensordict/register.py b/pandera/backends/tensordict/register.py index e69de29bb..77b4b85bd 100644 --- a/pandera/backends/tensordict/register.py +++ b/pandera/backends/tensordict/register.py @@ -0,0 +1,36 @@ +"""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 + from pandera.backends.tensordict.base import TensorDictSchemaBackend, _is_tensordict + from pandera.backends.tensordict.checks import TensorDictCheckBackend + + # Import builtin checks to register check functions + from pandera.backends.tensordict import builtin_checks + + 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/tensordict.py b/pandera/tensordict.py index a3c4a1fd1..a21ad3d01 100644 --- a/pandera/tensordict.py +++ b/pandera/tensordict.py @@ -14,11 +14,10 @@ except ImportError: DataType = None # type: ignore[misc, assignment] -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__ = [ diff --git a/tests/tensordict/test_tensordict_container.py b/tests/tensordict/test_tensordict_container.py index 2805c7979..a0dc9f741 100644 --- a/tests/tensordict/test_tensordict_container.py +++ b/tests/tensordict/test_tensordict_container.py @@ -66,8 +66,8 @@ def test_tensordict_schema_creation(self): }, batch_size=(32,), ) - assert "observation" in schema.columns - assert "action" in schema.columns + assert "observation" in schema.keys + assert "action" in schema.keys assert schema.batch_size == (32,) def test_tensordict_schema_from_list(self): @@ -77,8 +77,8 @@ def test_tensordict_schema_from_list(self): schema = TensorDictSchema( keys=["observation", "action"], batch_size=(32,) ) - assert "observation" in schema.columns - assert "action" in schema.columns + assert "observation" in schema.keys + assert "action" in schema.keys def test_tensordict_schema_repr(self): """Test TensorDictSchema repr.""" From f85c9c62e7f124cf8b683d9bad3c0033075bb110 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 20 Apr 2026 10:52:15 -0400 Subject: [PATCH 14/27] Fix mypy errors in TensorDict backend - Fixed preprocess() signature to include key parameter - Changed CoreCheckResult to CheckResult for error handling - Fixed register_builtin_check fallback signature - Added missing imports (BaseSchemaBackend, BaseConfig, Tensor) - Fixed annotation access from .annotation to .raw_annotation --- pandera/api/tensordict/components.py | 10 +-- pandera/api/tensordict/container.py | 47 ++++++++------ pandera/api/tensordict/model.py | 63 +++++++++---------- pandera/backends/tensordict/base.py | 34 +++++----- pandera/backends/tensordict/builtin_checks.py | 35 +++++++++-- pandera/backends/tensordict/checks.py | 30 ++++----- pandera/backends/tensordict/register.py | 21 ++++--- 7 files changed, 138 insertions(+), 102 deletions(-) diff --git a/pandera/api/tensordict/components.py b/pandera/api/tensordict/components.py index ba967bfd3..d793ce85a 100644 --- a/pandera/api/tensordict/components.py +++ b/pandera/api/tensordict/components.py @@ -10,7 +10,7 @@ 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. """ @@ -18,11 +18,11 @@ class Tensor: def __init__( self, dtype: Any = None, - shape: tuple[Optional[int], ...] | None = None, + shape: tuple[int | None, ...] | None = None, checks: CheckList | None = None, nullable: bool = False, coerce: bool = False, - name: Optional[str] = None, + name: str | None = None, required: bool = True, ) -> None: """Create tensor validator. @@ -47,7 +47,7 @@ def __init__( else None ) self.shape = shape - + # Normalize checks to a list self.checks: list = [] if checks is not None: @@ -55,7 +55,7 @@ def __init__( self.checks = list(checks) else: self.checks = [checks] - + self.nullable = nullable self.coerce = coerce self.name = name diff --git a/pandera/api/tensordict/container.py b/pandera/api/tensordict/container.py index 39cc1e957..ea260cc9c 100644 --- a/pandera/api/tensordict/container.py +++ b/pandera/api/tensordict/container.py @@ -6,29 +6,30 @@ 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: Optional[dict[str, Tensor]] = None, + keys: dict[str, Tensor] | None = None, batch_size: tuple[int | None, ...] | None = None, dtype: Any = None, checks=None, coerce: bool = False, - name: Optional[str] = None, - title: Optional[str] = None, - description: Optional[str] = None, - metadata: Optional[dict] = None, + name: str | None = None, + title: str | None = None, + description: str | None = None, + metadata: dict | None = None, ) -> None: """Create TensorDict schema. @@ -50,7 +51,7 @@ def __init__( keys = {} elif isinstance(keys, list): keys = {k: Tensor() for k in keys} - + super().__init__( dtype=dtype, checks=checks, @@ -60,22 +61,25 @@ def __init__( 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) -> BaseSchemaBackend: + 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. @@ -84,33 +88,36 @@ def get_backend(cls, check_obj: Any | None = None) -> BaseSchemaBackend: 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) + 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'): + + 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, diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index 10addacbc..49748dabb 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -2,29 +2,31 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, ClassVar, Optional +from typing import Any, ClassVar from pandera.api.base.model import BaseModel +from pandera.api.tensordict.components import Tensor +from pandera.api.tensordict.model_config import BaseConfig class TensorDictModel(BaseModel): """Declarative TensorDict schema using class definitions. - + Example: >>> import torch >>> import pandera.tensordict as pa - + >>> class MySchema(pa.TensorDictModel): ... observation: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 10)) ... action: pa.DataType = pa.Field(dtype=torch.int64, shape=(None,)) ... ... class Config: ... batch_size = (32,) - + >>> schema = MySchema.to_schema() """ - Config: type[object] = object + Config: type[BaseConfig] = BaseConfig __schema__: ClassVar[Any | None] = None def __init_subclass__(cls, **kwargs: Any) -> None: @@ -39,52 +41,49 @@ def to_schema(cls: type[TensorDictModel]) -> Any: :returns: TensorDictSchema with fields converted to Tensor components. """ columns: dict[str, Any] = {} - batch_size: Optional[tuple[int | None, ...]] = None + batch_size: tuple[int | None, ...] | None = None # Get field definitions from class annotations and Field info for name, (annotation_info, field_info) in cls.__fields__.items(): if hasattr(field_info, "shape"): # Extract dtype from annotation - dtype = annotation_info.annotation - + dtype = annotation_info.raw_annotation + # 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): + 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 = getattr(field_info, "shape", None) - - # Build Tensor component - use Tensor directly since it's not imported yet - try: - from pandera.tensordict import Tensor - except ImportError: - Tensor = None - - if Tensor is not None: - tensor_kwargs: dict[str, Any] = { - "dtype": dtype, - "shape": shape, - "name": name, - } - - if hasattr(field_info, "checks") and field_info.checks: - tensor_kwargs["checks"] = field_info.checks - - columns[name] = Tensor(**tensor_kwargs) - + + tensor_kwargs: dict[str, Any] = { + "dtype": dtype, + "shape": shape, + "name": name, + } + + if hasattr(field_info, "checks") and field_info.checks: + tensor_kwargs["checks"] = field_info.checks + + columns[name] = Tensor(**tensor_kwargs) + # Get batch_size from Config if hasattr(cls, "Config"): batch_size = getattr(cls.Config, "batch_size", None) try: from pandera.tensordict import TensorDictSchema - + return TensorDictSchema(keys=columns, batch_size=batch_size) except ImportError: return None diff --git a/pandera/backends/tensordict/base.py b/pandera/backends/tensordict/base.py index a3228ea98..6c37ef3e9 100644 --- a/pandera/backends/tensordict/base.py +++ b/pandera/backends/tensordict/base.py @@ -21,16 +21,16 @@ def _is_tensordict(obj: Any) -> bool: 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: + if hasattr(obj, "_is_tensorclass") and obj._is_tensorclass: return True - + return False except ImportError: return False @@ -43,11 +43,11 @@ def _get_tensor(obj: Any, key: str) -> Any: 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") @@ -156,22 +156,22 @@ def _check_keys(self, check_obj, schema, error_handler, lazy): 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'): + 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( @@ -291,15 +291,19 @@ def _run_value_checks(self, check_obj, schema, error_handler, lazy): for check_index, check in enumerate(tensor_schema.checks): try: # Use TensorDictCheckBackend directly - from pandera.backends.tensordict.checks import TensorDictCheckBackend - + 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) + check_result = check_backend.postprocess( + tensor, check_result + ) except Exception as exc: check_result = CoreCheckResult( passed=False, diff --git a/pandera/backends/tensordict/builtin_checks.py b/pandera/backends/tensordict/builtin_checks.py index 53a12becb..c89f55506 100644 --- a/pandera/backends/tensordict/builtin_checks.py +++ b/pandera/backends/tensordict/builtin_checks.py @@ -5,17 +5,21 @@ # 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 typing import Any, Iterable +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})", @@ -32,7 +36,9 @@ def greater_than(data: torch.Tensor, min_value: Any) -> torch.Tensor: aliases=["ge"], error="greater_than_or_equal_to({min_value})", ) - def greater_than_or_equal_to(data: torch.Tensor, min_value: Any) -> torch.Tensor: + 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. @@ -56,7 +62,9 @@ def less_than(data: torch.Tensor, max_value: Any) -> torch.Tensor: aliases=["le"], error="less_than_or_equal_to({max_value})", ) - def less_than_or_equal_to(data: torch.Tensor, max_value: Any) -> torch.Tensor: + 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. @@ -106,7 +114,22 @@ def isin(data: torch.Tensor, allowed_values: Iterable) -> torch.Tensor: except ImportError: # extensions not available yet - def register_builtin_check(fn=None, **kwargs): + 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 - return decorator if fn is None else decorator(fn) + + 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 index 0ee529c18..2296020a5 100644 --- a/pandera/backends/tensordict/checks.py +++ b/pandera/backends/tensordict/checks.py @@ -22,19 +22,19 @@ def __init__(self, check: Check): else None ) - def preprocess(self, check_obj: Any) -> Any: + 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, @@ -47,21 +47,21 @@ def apply(self, check_obj: Any) -> CheckResult: try: result = self.check_fn(check_obj) except Exception as exc: - return CoreCheckResult( - passed=False, - check=self.check, - reason_code=SchemaErrorReason.CHECK_ERROR, - message=str(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, @@ -90,9 +90,7 @@ def _reduce_to_bool(self, result: Any) -> bool: # Assume it's already a boolean return bool(result) - def _get_failure_cases( - self, check_obj: Any, result: Any - ) -> Any | None: + def _get_failure_cases(self, check_obj: Any, result: Any) -> Any | None: """Get the failure cases from the check result.""" try: import torch @@ -107,7 +105,9 @@ def _get_failure_cases( return None - def postprocess(self, check_obj: Any, check_output: CheckResult) -> CheckResult: + 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 diff --git a/pandera/backends/tensordict/register.py b/pandera/backends/tensordict/register.py index 77b4b85bd..a59aa623c 100644 --- a/pandera/backends/tensordict/register.py +++ b/pandera/backends/tensordict/register.py @@ -6,31 +6,34 @@ @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 - from pandera.backends.tensordict.base import TensorDictSchemaBackend, _is_tensordict - from pandera.backends.tensordict.checks import TensorDictCheckBackend - + # 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 From 2fe0b006a1730f77676a2539c9a0705c59254944 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Mon, 20 Apr 2026 17:09:28 -0400 Subject: [PATCH 15/27] Fix TensorDictModel class-based implementation - Updated TensorDictModel to follow the same pattern as xarray.DatasetModel and DataFrameModel - Fields are now defined using pa.Field() directly in type annotations without needing custom _field classmethod - Type annotations specify dtype (e.g., torch.float32, torch.int64) instead of torch.Tensor - Added descriptors for field collection, schema caching, checks, and parsers - Exported SchemaError and SchemaErrors from pandera.tensordict module - Updated documentation to reflect correct API usage - Added comprehensive test coverage for TensorDictModel Fixes issue with class-based model implementation that required _field classmethod --- docs/source/pytorch_guide/index.md | 20 +- docs/source/pytorch_guide/tensordict_model.md | 89 ++- pandera/api/tensordict/model.py | 511 ++++++++++++++++-- pandera/api/tensordict/model_components.py | 1 + pandera/tensordict.py | 6 + specs/pytorch-tensordict.md | 140 ++++- tests/tensordict/test_tensordict_model.py | 126 +++++ 7 files changed, 812 insertions(+), 81 deletions(-) diff --git a/docs/source/pytorch_guide/index.md b/docs/source/pytorch_guide/index.md index 9c87dbf6c..be029935d 100644 --- a/docs/source/pytorch_guide/index.md +++ b/docs/source/pytorch_guide/index.md @@ -48,21 +48,27 @@ schema.validate(td) ```{code-cell} python class RL(pa.TensorDictModel): - observation: torch.Tensor - action: torch.Tensor + """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() - @classmethod - def _field(cls, field): - return field(dtype=torch.float32, shape=(None, 10)) + class Config: + batch_size = (32,) -# Validate using the model +# Validate using the model - schema is built automatically td = TensorDict( - {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + {"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} diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md index 26885a3b2..c4d66f109 100644 --- a/docs/source/pytorch_guide/tensordict_model.md +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -14,15 +14,19 @@ schema that maps directly to a {class}`~tensordict.TensorDict`. ```{code-cell} python import torch import pandera.tensordict as pa -from pandera import Check class RL(pa.TensorDictModel): - observation: pa.DataType - action: pa.DataType + """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 {class}`~pandera.tensordict.DataType` as the type annotation and -{func}`~pandera.tensordict.Field` to customize field options. +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 @@ -30,7 +34,7 @@ Use {class}`~pandera.tensordict.DataType` as the type annotation and from tensordict import TensorDict td = TensorDict( - {"observation": torch.randn(32, 10), "action": torch.randn(32, 5)}, + {"observation": torch.randn(32, 10), "action": torch.randint(0, 4, (32,))}, batch_size=[32], ) validated = RL.validate(td) @@ -40,21 +44,78 @@ validated = RL.validate(td) Use {func}`~pandera.tensordict.Field` to customize field options: -- `dtype`: Torch dtype -- `shape`: Expected shape tuple -- `checks`: List of Check instances +- `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 +from pandera import Check + class RLWithConfig(pa.TensorDictModel): - observation: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 10)) - action: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 5)) - reward: pa.DataType = pa.Field( - dtype=torch.float32, + """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,), - checks=[Check.greater_than(0.0)], + checks=[Check.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 diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index 49748dabb..b86a5a254 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -2,23 +2,153 @@ from __future__ import annotations -from typing import Any, ClassVar +import inspect +import threading +import typing +from typing import Any, ClassVar, cast, get_type_hints + +import typing_inspect 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 + + +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] + + +DATAFRAME_CHECK_KEY = "__dataframe_check_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: pa.DataType = pa.Field(dtype=torch.float32, shape=(None, 10)) - ... action: pa.DataType = pa.Field(dtype=torch.int64, shape=(None,)) + ... observation: torch.float32 = pa.Field(shape=(None, 10)) + ... action: torch.int64 = pa.Field(shape=(None,)) ... ... class Config: ... batch_size = (32,) @@ -27,66 +157,349 @@ class TensorDictModel(BaseModel): """ Config: type[BaseConfig] = BaseConfig - __schema__: ClassVar[Any | None] = None + __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) - cls.__schema__ = cls.to_schema() + hints = get_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_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, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.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 to_schema(cls: type[TensorDictModel]) -> Any: + 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, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.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, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.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] = [] + checks[field].append(check_) + return checks + + @classmethod + def _extract_df_checks(cls, check_infos: list[Any]) -> list[Check]: + """Collect dataframe-level checks.""" + return [ + check_info.to_check(cls) + for check_info in check_infos + if hasattr(check_info, "to_check") + ] + + @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, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.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] = [] + parsers[field].append(parser_) + return parsers + + @classmethod + def _extract_df_parsers(cls, parser_infos: list[Any]) -> list[Parser]: + """Collect dataframe-level parsers.""" + return [ + parser_info.to_parser(cls) + for parser_info in parser_infos + if hasattr(parser_info, "to_parser") + ] + + @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] = {} - batch_size: tuple[int | None, ...] | None = None - - # Get field definitions from class annotations and Field info - for name, (annotation_info, field_info) in cls.__fields__.items(): - if hasattr(field_info, "shape"): - # Extract dtype from annotation - dtype = annotation_info.raw_annotation - - # 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 = getattr(field_info, "shape", None) - - tensor_kwargs: dict[str, Any] = { - "dtype": dtype, - "shape": shape, - "name": name, - } - - if hasattr(field_info, "checks") and field_info.checks: - tensor_kwargs["checks"] = field_info.checks - - columns[name] = Tensor(**tensor_kwargs) + + 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 - if hasattr(cls, "Config"): - batch_size = getattr(cls.Config, "batch_size", None) + batch_size: tuple[int | None, ...] | None = ( + getattr(cfg, "batch_size", None) if cfg else None + ) try: - from pandera.tensordict import TensorDictSchema + return TensorDictSchema(keys=columns or None, batch_size=batch_size) + except ImportError as e: + raise RuntimeError( + "Could not import TensorDictSchema" + ) from e - return TensorDictSchema(keys=columns, batch_size=batch_size) - except ImportError: - return None + @classmethod + def to_schema(cls: type[TensorDictModel]) -> TensorDictSchema: + """Create :class:`~pandera.TensorDictSchema` from the :class:`.TensorDictModel`.""" + return cls.__schema__ @classmethod def validate( @@ -115,3 +528,7 @@ def validate( 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 index afc70b88c..b83850740 100644 --- a/pandera/api/tensordict/model_components.py +++ b/pandera/api/tensordict/model_components.py @@ -9,6 +9,7 @@ from pandera.errors import SchemaInitError CHECK_KEY = "__check_config__" +TENSOR_CHECK_KEY = "__tensor_check_config__" class TensorDictFieldInfo(FieldInfo): diff --git a/pandera/tensordict.py b/pandera/tensordict.py index a21ad3d01..65380b0b5 100644 --- a/pandera/tensordict.py +++ b/pandera/tensordict.py @@ -27,4 +27,10 @@ "Field", "errors", "DataType", + # Error classes + "SchemaError", + "SchemaErrors", ] + +# Import error classes for convenience +from pandera.errors import SchemaError, SchemaErrors diff --git a/specs/pytorch-tensordict.md b/specs/pytorch-tensordict.md index b1620e062..fe0765b53 100644 --- a/specs/pytorch-tensordict.md +++ b/specs/pytorch-tensordict.md @@ -158,7 +158,7 @@ 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 @@ -168,14 +168,13 @@ from pandera import Check class PolicyModel(pa.TensorDictModel): """Schema for a policy network output.""" - logits: torch.Tensor = pa.Field( - dtype=torch.float32, + # Dtype is specified in the type annotation, Field defines constraints + logits: torch.float32 = pa.Field( shape=(None, 10), ge=-5.0, le=5.0 ) - value: torch.Tensor = pa.Field( - dtype=torch.float32, + value: torch.float32 = pa.Field( shape=(None,), gt=0.0 ) @@ -183,14 +182,8 @@ class PolicyModel(pa.TensorDictModel): 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 @@ -841,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 diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index b8035e5ec..df370b0d4 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -17,3 +17,129 @@ 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 TensorDictModel, Field + + 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 TensorDictModel, Field + + 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 TensorDictModel, Field + + 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 TensorDictModel, Field + + 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 TensorDictModel, Field + + 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"]) From 8b72b9356a5dcf18ad43f723d158193347a9908b Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 21 Apr 2026 09:10:05 -0400 Subject: [PATCH 16/27] Fix pre-commit hook failures - Add missing CHECK_KEY and PARSER_KEY constants to model.py - Fix mypy type errors by filtering None values from extract methods - Apply ruff formatting fixes --- docs/source/pytorch_guide/index.md | 2 +- docs/source/pytorch_guide/tensordict_model.md | 4 +- pandera/api/tensordict/model.py | 134 ++++++++++++------ specs/pytorch-tensordict.md | 10 +- tests/tensordict/test_tensordict_model.py | 10 +- 5 files changed, 101 insertions(+), 59 deletions(-) diff --git a/docs/source/pytorch_guide/index.md b/docs/source/pytorch_guide/index.md index be029935d..d16c1e093 100644 --- a/docs/source/pytorch_guide/index.md +++ b/docs/source/pytorch_guide/index.md @@ -49,7 +49,7 @@ schema.validate(td) ```{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,)) diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md index c4d66f109..8d5f79704 100644 --- a/docs/source/pytorch_guide/tensordict_model.md +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -17,7 +17,7 @@ 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,)) @@ -54,7 +54,7 @@ from pandera import Check class RLWithConfig(pa.TensorDictModel): """RL schema with field-level checks.""" - + observation: torch.float32 = pa.Field( shape=(None, 10), ge=-1.0, diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index b86a5a254..d0e29c319 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -23,7 +23,7 @@ TParsers = dict[str, list[Parser]] _CONFIG_KEY = "Config" -MODEL_CACHE: dict[tuple[type["TensorDictModel"], int], Any] = {} +MODEL_CACHE: dict[tuple[type[TensorDictModel], int], Any] = {} def _is_field(name: str) -> bool: @@ -112,7 +112,9 @@ def __get__(self, obj: Any, cls: type[TensorDictModel]) -> list[Parser]: return self.cache[cls] +CHECK_KEY = "__check_config__" DATAFRAME_CHECK_KEY = "__dataframe_check_config__" +PARSER_KEY = "__parser_config__" DATAFRAME_PARSER_KEY = "__dataframe_parser_config__" @@ -139,21 +141,21 @@ def _convert_extras_to_checks(extras: dict[str, Any]) -> list[Check]: 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: + 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 + 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,) + >>> 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() + >>> schema = MySchema.to_schema() """ Config: type[BaseConfig] = BaseConfig @@ -174,10 +176,10 @@ def __init_subclass__(cls, **kwargs: Any) -> None: cls.Config.name = None else: cls.Config = type("Config", (cls.Config,), {}) - + super().__init_subclass__(**kwargs) hints = get_type_hints(cls, include_extras=True) - + for fname in hints.keys(): if not _is_field(fname): continue @@ -185,11 +187,11 @@ def __init_subclass__(cls, **kwargs: Any) -> None: 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 @@ -209,7 +211,7 @@ def _get_model_attrs(cls) -> dict[str, Any]: def _collect_fields(cls) -> TFields: """Centralize publicly named fields and their corresponding annotations.""" from pandera.api.tensordict.model_components import TensorDictFieldInfo - + annotations = get_type_hints(cls, include_extras=True) attrs = cls._get_model_attrs() @@ -230,7 +232,7 @@ def _collect_fields(cls) -> TFields: 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( @@ -246,11 +248,13 @@ def _extract_config_options_and_extras( 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 - + 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( @@ -310,7 +314,14 @@ def _collect_check_infos(cls, key: str) -> list[Any]: 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, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.get("Field", object), "__bases__") else object): + if not isinstance( + check_info, + type(cls).__dict__.get("Field", object).__bases__[0] + if hasattr( + type(cls).__dict__.get("Field", object), "__bases__" + ) + else object, + ): continue if attr_name in method_names: continue @@ -331,7 +342,14 @@ def _collect_parser_infos(cls, key: str) -> list[Any]: 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, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.get("Field", object), "__bases__") else object): + if not isinstance( + parser_info, + type(cls).__dict__.get("Field", object).__bases__[0] + if hasattr( + type(cls).__dict__.get("Field", object), "__bases__" + ) + else object, + ): continue method_names.add(attr_name) parser_infos.append(parser_info) @@ -357,10 +375,17 @@ def _extract_checks( for check_info in check_infos: if not hasattr(check_info, "fields"): continue - + check_info_fields = { field.name - if isinstance(field, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.get("Field", object), "__bases__") else object) + if isinstance( + field, + type(cls).__dict__.get("Field", object).__bases__[0] + if hasattr( + type(cls).__dict__.get("Field", object), "__bases__" + ) + else object, + ) else field for field in check_info.fields } @@ -372,7 +397,9 @@ def _extract_checks( ) check_ = ( - check_info.to_check(cls) if hasattr(check_info, "to_check") else None + check_info.to_check(cls) + if hasattr(check_info, "to_check") + else None ) for field in matched: @@ -382,17 +409,20 @@ def _extract_checks( ) if field not in checks: checks[field] = [] - checks[field].append(check_) + 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.""" - return [ - check_info.to_check(cls) - for check_info in check_infos - if hasattr(check_info, "to_check") - ] + 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( @@ -403,10 +433,17 @@ def _extract_parsers( for parser_info in parser_infos: if not hasattr(parser_info, "fields"): continue - + parser_info_fields = { field.name - if isinstance(field, type(cls).__dict__.get("Field", object).__bases__[0] if hasattr(type(cls).__dict__.get("Field", object), "__bases__") else object) + if isinstance( + field, + type(cls).__dict__.get("Field", object).__bases__[0] + if hasattr( + type(cls).__dict__.get("Field", object), "__bases__" + ) + else object, + ) else field for field in parser_info.fields } @@ -430,17 +467,20 @@ def _extract_parsers( ) if field not in parsers: parsers[field] = [] - parsers[field].append(parser_) + 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.""" - return [ - parser_info.to_parser(cls) - for parser_info in parser_infos - if hasattr(parser_info, "to_parser") - ] + 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: @@ -464,7 +504,9 @@ def build_schema_(cls, **kwargs) -> TensorDictSchema: DataType as _DataType, ) - if isinstance(dtype, type) and issubclass(dtype, _DataType): + if isinstance(dtype, type) and issubclass( + dtype, _DataType + ): dtype = dtype.type except Exception: pass @@ -490,11 +532,11 @@ def build_schema_(cls, **kwargs) -> TensorDictSchema: ) try: - return TensorDictSchema(keys=columns or None, batch_size=batch_size) + return TensorDictSchema( + keys=columns or None, batch_size=batch_size + ) except ImportError as e: - raise RuntimeError( - "Could not import TensorDictSchema" - ) from e + raise RuntimeError("Could not import TensorDictSchema") from e @classmethod def to_schema(cls: type[TensorDictModel]) -> TensorDictSchema: diff --git a/specs/pytorch-tensordict.md b/specs/pytorch-tensordict.md index fe0765b53..28b8f9c33 100644 --- a/specs/pytorch-tensordict.md +++ b/specs/pytorch-tensordict.md @@ -854,7 +854,7 @@ 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( @@ -862,10 +862,10 @@ class RLModel(pa.TensorDictModel): checks=Check.isin([0, 1, 2, 3]) ) reward: torch.float32 = pa.Field() - + class Config: batch_size = (32,) - + # Access schema directly schema = RLModel.to_schema() @@ -884,7 +884,7 @@ validated = RLModel.validate(batch) ```python class ActorCriticOutput(pa.TensorDictModel): """Actor-critic network output.""" - + action_mean: torch.float32 = pa.Field( shape=(None, 4), ge=-2.0, @@ -943,7 +943,7 @@ For reusable, self-documenting schemas with type hints, use `TensorDictModel`: class MySchema(pa.TensorDictModel): observation: torch.float32 = pa.Field(shape=(None, 10)) action: torch.int64 = pa.Field(shape=(None,)) - + class Config: batch_size = (32,) diff --git a/tests/tensordict/test_tensordict_model.py b/tests/tensordict/test_tensordict_model.py index df370b0d4..c0c5da916 100644 --- a/tests/tensordict/test_tensordict_model.py +++ b/tests/tensordict/test_tensordict_model.py @@ -25,7 +25,7 @@ class TestTensorDictModelBasic: def test_model_with_field_annotations(self): """Test model with field annotations.""" - from pandera.tensordict import TensorDictModel, Field + from pandera.tensordict import Field, TensorDictModel class MyModel(TensorDictModel): observation: torch.float32 = Field(shape=(None, 10)) @@ -40,7 +40,7 @@ class Config: def test_model_validation(self): """Test model validation with valid data.""" - from pandera.tensordict import TensorDictModel, Field + from pandera.tensordict import Field, TensorDictModel class MyModel(TensorDictModel): observation: torch.float32 = Field(shape=(None, 10)) @@ -59,7 +59,7 @@ class Config: def test_model_validation_failure(self): """Test model validation fails with invalid data.""" import pandera.errors as errors - from pandera.tensordict import TensorDictModel, Field + from pandera.tensordict import Field, TensorDictModel class MyModel(TensorDictModel): observation: torch.float32 = Field(shape=(None, 10)) @@ -84,7 +84,7 @@ 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 TensorDictModel, Field + from pandera.tensordict import Field, TensorDictModel class MyModel(TensorDictModel): values: torch.float32 = Field( @@ -120,7 +120,7 @@ class TestTensorDictModelInheritance: def test_model_inheritance(self): """Test model can be inherited from.""" - from pandera.tensordict import TensorDictModel, Field + from pandera.tensordict import Field, TensorDictModel class BaseModel(TensorDictModel): observation: torch.float32 = Field(shape=(None, 10)) From cb9d015c46f2c242bcc3be31ba89c35777b94034 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 21 Apr 2026 14:50:30 -0400 Subject: [PATCH 17/27] Fix PyTorch documentation examples - error_reporting.md: Use str(err) instead of non-existent err.message attribute - tensordict_checks.md: Fix test data to pass less_than(1.0) check (0.95 vs 1.0) - tensordict_model.md: Fix indentation and replace unsupported checks parameter with isin - tensordict_schema.md: Update example to show validation errors without TensorDict creation conflicts --- docs/source/pytorch_guide/error_reporting.md | 2 +- docs/source/pytorch_guide/tensordict_checks.md | 2 +- docs/source/pytorch_guide/tensordict_model.md | 6 ++---- docs/source/pytorch_guide/tensordict_schema.md | 15 +++++++++------ 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/docs/source/pytorch_guide/error_reporting.md b/docs/source/pytorch_guide/error_reporting.md index ed9e01289..9d5fac08e 100644 --- a/docs/source/pytorch_guide/error_reporting.md +++ b/docs/source/pytorch_guide/error_reporting.md @@ -47,7 +47,7 @@ try: except pa.errors.SchemaErrors as exc: print(f"Found {len(exc.schema_errors)} errors") for err in exc.schema_errors: - print(f" - {err.message}") + print(f" - {str(err)}") print(f"\nError counts: {exc.error_counts}") ``` diff --git a/docs/source/pytorch_guide/tensordict_checks.md b/docs/source/pytorch_guide/tensordict_checks.md index a1e8f80ae..f06e6bd5b 100644 --- a/docs/source/pytorch_guide/tensordict_checks.md +++ b/docs/source/pytorch_guide/tensordict_checks.md @@ -28,7 +28,7 @@ schema = pa.TensorDictSchema( ) td = TensorDict( - {"values": torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])}, + {"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) diff --git a/docs/source/pytorch_guide/tensordict_model.md b/docs/source/pytorch_guide/tensordict_model.md index 8d5f79704..4c0c93234 100644 --- a/docs/source/pytorch_guide/tensordict_model.md +++ b/docs/source/pytorch_guide/tensordict_model.md @@ -34,7 +34,7 @@ define additional constraints. from tensordict import TensorDict td = TensorDict( - {"observation": torch.randn(32, 10), "action": torch.randint(0, 4, (32,))}, + {"observation": torch.randn(32, 10), "action": torch.randint(0, 4, (32,)), "reward": torch.randn(32)}, batch_size=[32], ) validated = RL.validate(td) @@ -50,8 +50,6 @@ Use {func}`~pandera.tensordict.Field` to customize field options: - `default`: Default value if missing ```{code-cell} python -from pandera import Check - class RLWithConfig(pa.TensorDictModel): """RL schema with field-level checks.""" @@ -62,7 +60,7 @@ class RLWithConfig(pa.TensorDictModel): ) action: torch.int64 = pa.Field( shape=(None,), - checks=[Check.isin([0, 1, 2, 3])], + isin=[0, 1, 2, 3], ) reward: torch.float32 = pa.Field( gt=0.0, diff --git a/docs/source/pytorch_guide/tensordict_schema.md b/docs/source/pytorch_guide/tensordict_schema.md index 7362bc172..56eaf8b5d 100644 --- a/docs/source/pytorch_guide/tensordict_schema.md +++ b/docs/source/pytorch_guide/tensordict_schema.md @@ -43,18 +43,21 @@ validated = schema.validate(td) ## Lazy validation -Set `lazy=True` to collect all errors instead of fail-fast: +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_invalid = TensorDict( - {"observation": torch.randn(16, 10), "action": torch.randn(32, 5)}, - batch_size=[16], +td_wrong_dtype = TensorDict( + {"observation": torch.randn(32, 10), "action": torch.randint(0, 5, (32,))}, + batch_size=[32], ) try: - schema.validate(td_invalid, lazy=True) + schema.validate(td_wrong_dtype, lazy=True) except pa.errors.SchemaErrors as exc: - print(exc) + print(f"Total errors: {len(exc.schema_errors)}") + for err in exc.schema_errors: + print(f" - {str(err)}") ``` ## Tensor component From df80fecc73d95634e84d44bc81886020093cdf07 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 21 Apr 2026 15:22:15 -0400 Subject: [PATCH 18/27] Add torch and tensordict to docs dependency group --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1467167a4..5676962c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,6 +170,8 @@ docs = [ "sphinx-docsearch", "grpcio", "ray", + "torch", + "tensordict", "types-click", "types-pytz", "types-pyyaml", From 23dc6dc59061bb7f2f83ac28112de6b4ee2e9ac8 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 21 Apr 2026 16:24:35 -0400 Subject: [PATCH 19/27] Remove torch from docs group (already in all) --- pyproject.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5676962c5..1467167a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,8 +170,6 @@ docs = [ "sphinx-docsearch", "grpcio", "ray", - "torch", - "tensordict", "types-click", "types-pytz", "types-pyyaml", From 76ebcd88db14369366ed91a4cb057da609115343 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 21 Apr 2026 16:25:57 -0400 Subject: [PATCH 20/27] Install all extras in docs session to ensure torch is available --- noxfile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index 129a817a9..68b569f2b 100644 --- a/noxfile.py +++ b/noxfile.py @@ -315,7 +315,7 @@ 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("-e", ".[all]") session.install( *_testing_requirements( session, extra="all", pandas=PANDAS_VERSIONS[0] From d6b1ea27a64585081e2eefd70128dc2cf8493c7c Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 21 Apr 2026 21:57:42 -0400 Subject: [PATCH 21/27] Re-add torch and tensordict to docs dependency group --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1467167a4..5676962c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,6 +170,8 @@ docs = [ "sphinx-docsearch", "grpcio", "ray", + "torch", + "tensordict", "types-click", "types-pytz", "types-pyyaml", From 7df739881d54202fa5d6980ba27921b8fb5bf66b Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Wed, 22 Apr 2026 15:16:22 -0400 Subject: [PATCH 22/27] Fix: ensure torch is available for xdoctest in TensorDictModel - Add optional torch import to model.py module namespace - Remove redundant dependency installation in docs session --- noxfile.py | 11 ++++------- pandera/api/tensordict/model.py | 6 ++++++ 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/noxfile.py b/noxfile.py index 68b569f2b..548b56faa 100644 --- a/noxfile.py +++ b/noxfile.py @@ -316,12 +316,6 @@ def docs(session: Session) -> None: # this is needed until ray and geopandas are supported on python 3.10 session.install("-e", ".[all]") - session.install( - *_testing_requirements( - session, extra="all", pandas=PANDAS_VERSIONS[0] - ), - *nox.project.dependency_groups(PYPROJECT, "dev", "testing", "docs"), - ) session.run("uv", "pip", "list") session.chdir("docs") @@ -357,7 +351,10 @@ def docs(session: Session) -> None: "sphinx-build", *args, ) - + + # Ensure torch is available for TensorDictModel doctests + session.run("python", "-c", "import torch") + session.run("xdoctest", PACKAGE, "--quiet") diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index d0e29c319..b1b7a1850 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -9,6 +9,12 @@ 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 0e6a650bd66f9906e6f40a6688604f91aba161b4 Mon Sep 17 00:00:00 2001 From: cosmicBboy Date: Thu, 23 Apr 2026 16:16:56 -0400 Subject: [PATCH 23/27] fix docs tests Signed-off-by: cosmicBboy --- noxfile.py | 10 +++++++--- pyproject.toml | 1 + 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/noxfile.py b/noxfile.py index 548b56faa..e0ec482b0 100644 --- a/noxfile.py +++ b/noxfile.py @@ -315,7 +315,11 @@ def docs(session: Session) -> None: """Build the documentation.""" # this is needed until ray and geopandas are supported on python 3.10 - session.install("-e", ".[all]") + session.install( + "-e", + ".[all]", + *nox.project.dependency_groups(PYPROJECT, "docs"), + ) session.run("uv", "pip", "list") session.chdir("docs") @@ -351,10 +355,10 @@ def docs(session: Session) -> None: "sphinx-build", *args, ) - + # Ensure torch is available for TensorDictModel doctests session.run("python", "-c", "import torch") - + session.run("xdoctest", PACKAGE, "--quiet") diff --git a/pyproject.toml b/pyproject.toml index 5676962c5..aa648b6f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,6 +177,7 @@ docs = [ "types-pyyaml", "types-requests", "types-setuptools", + "xdoctest", ] [tool.setuptools] From 6224a542dec83cd7d19f1695281257ef2361c5ad Mon Sep 17 00:00:00 2001 From: cosmicBboy Date: Thu, 23 Apr 2026 16:40:51 -0400 Subject: [PATCH 24/27] updates Signed-off-by: cosmicBboy --- pandera/api/polars/model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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) From b727fffb4de3c1b481acb2df9214521791c42ce2 Mon Sep 17 00:00:00 2001 From: cosmicBboy Date: Thu, 23 Apr 2026 20:20:04 -0400 Subject: [PATCH 25/27] update ibis docs Signed-off-by: cosmicBboy --- docs/source/ibis.md | 6 +++++- noxfile.py | 4 ---- 2 files changed, 5 insertions(+), 5 deletions(-) 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/noxfile.py b/noxfile.py index e0ec482b0..9d0038758 100644 --- a/noxfile.py +++ b/noxfile.py @@ -177,10 +177,6 @@ def _testing_requirements( req = "ibis-framework[duckdb] >= 11.0.0" if req == "polars": req = f"polars=={polars}" - if req == "torch": - req = "torch" - if req == "tensordict": - req = "tensordict" # for some reason uv will try to install an old version of dask, # have to specifically pin dask[dataframe] to a higher version From a343ef7a71fb655e0ea98172053df099b0c1e2af Mon Sep 17 00:00:00 2001 From: cosmicBboy Date: Thu, 23 Apr 2026 20:36:48 -0400 Subject: [PATCH 26/27] update Signed-off-by: cosmicBboy --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index aa648b6f2..a7b42c1ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -159,6 +159,7 @@ testing = [ ] docs = [ "setuptools", + "duckdb", "pyarrow-hotfix", "sphinx<9", "sphinx-design", From 8550274b959b3872559a6fa5a828618b5e7f566e Mon Sep 17 00:00:00 2001 From: cosmicBboy Date: Thu, 23 Apr 2026 21:05:49 -0400 Subject: [PATCH 27/27] update Signed-off-by: cosmicBboy --- pandera/api/tensordict/model.py | 86 ++++++++++++++++++++++++++++----- 1 file changed, 75 insertions(+), 11 deletions(-) diff --git a/pandera/api/tensordict/model.py b/pandera/api/tensordict/model.py index b1b7a1850..dadd6ced9 100644 --- a/pandera/api/tensordict/model.py +++ b/pandera/api/tensordict/model.py @@ -3,9 +3,11 @@ from __future__ import annotations import inspect +import sys import threading +import types as types_module import typing -from typing import Any, ClassVar, cast, get_type_hints +from typing import Any, ClassVar, cast import typing_inspect @@ -36,6 +38,62 @@ 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] = {} @@ -184,7 +242,7 @@ def __init_subclass__(cls, **kwargs: Any) -> None: cls.Config = type("Config", (cls.Config,), {}) super().__init_subclass__(**kwargs) - hints = get_type_hints(cls, include_extras=True) + hints = _get_tensor_dict_class_type_hints(cls, include_extras=True) for fname in hints.keys(): if not _is_field(fname): @@ -218,7 +276,9 @@ def _collect_fields(cls) -> TFields: """Centralize publicly named fields and their corresponding annotations.""" from pandera.api.tensordict.model_components import TensorDictFieldInfo - annotations = get_type_hints(cls, include_extras=True) + annotations = _get_tensor_dict_class_type_hints( + cls, include_extras=True + ) attrs = cls._get_model_attrs() missing = [] @@ -322,9 +382,10 @@ def _collect_check_infos(cls, key: str) -> list[Any]: check_info = getattr(attr_value, key, None) if not isinstance( check_info, - type(cls).__dict__.get("Field", object).__bases__[0] + _metaclass_dict(cls).get("Field", object).__bases__[0] if hasattr( - type(cls).__dict__.get("Field", object), "__bases__" + _metaclass_dict(cls).get("Field", object), + "__bases__", ) else object, ): @@ -350,9 +411,10 @@ def _collect_parser_infos(cls, key: str) -> list[Any]: parser_info = getattr(attr_value, key, None) if not isinstance( parser_info, - type(cls).__dict__.get("Field", object).__bases__[0] + _metaclass_dict(cls).get("Field", object).__bases__[0] if hasattr( - type(cls).__dict__.get("Field", object), "__bases__" + _metaclass_dict(cls).get("Field", object), + "__bases__", ) else object, ): @@ -386,9 +448,10 @@ def _extract_checks( field.name if isinstance( field, - type(cls).__dict__.get("Field", object).__bases__[0] + _metaclass_dict(cls).get("Field", object).__bases__[0] if hasattr( - type(cls).__dict__.get("Field", object), "__bases__" + _metaclass_dict(cls).get("Field", object), + "__bases__", ) else object, ) @@ -444,9 +507,10 @@ def _extract_parsers( field.name if isinstance( field, - type(cls).__dict__.get("Field", object).__bases__[0] + _metaclass_dict(cls).get("Field", object).__bases__[0] if hasattr( - type(cls).__dict__.get("Field", object), "__bases__" + _metaclass_dict(cls).get("Field", object), + "__bases__", ) else object, )