feat: PyTorch TensorDict Phase 4 - Schema Inference, IO Serialization, and Hypothesis Strategies

docs/source/pytorch_guide/index.md

@@ -44,6 +44,30 @@ td = TensorDict(
 schema.validate(td)
 ```
 
+## Type Coercion
+
+Set `coerce=True` to automatically convert tensor dtypes:
+
+```{code-cell} python
+schema_coerce = pa.TensorDictSchema(
+    keys={
+        "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)),
+    },
+    batch_size=(32,),
+    coerce=True,
+)
+
+# Input with wrong dtype (float64)
+td_wrong_dtype = TensorDict(
+    {"observation": torch.randn(32, 10).to(torch.float64)},
+    batch_size=[32],
+)
+
+# Dtype is automatically coerced to float32
+validated = schema_coerce.validate(td_wrong_dtype)
+assert validated["observation"].dtype == torch.float32
+```
+
 ## Define a schema with a class-based model
 
 ```{code-cell} python
@@ -78,12 +102,18 @@ Use {func}`~pandera.tensordict.Field` to define additional constraints like shap
 tensordict_schema
 tensordict_model
 tensordict_checks
+tensordict_schema_inference
+tensordict_io
+tensordict_strategies
 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-tensordict-inference` — infer schemas from data automatically
+- {ref}`pytorch-tensordict-io` — save/load schemas with YAML/JSON
+- {ref}`pytorch-tensordict-strategies` — generate synthetic data with Hypothesis
 - {ref}`pytorch-error-reporting` — `SchemaError` / `SchemaErrors`, lazy validation, and failure cases
 
 ## See also

docs/source/pytorch_guide/tensordict_io.md

@@ -0,0 +1,227 @@
+---
+file_format: mystnb
+---
+
+(pytorch-tensordict-io)=
+
+# Serialization and Deserialization
+
+Save and load TensorDict schemas using YAML or JSON, or save TensorDicts with embedded schema metadata.
+
+## Schema Serialization
+
+### YAML Format
+
+```{code-cell} python
+import torch
+from pandera import Check
+import pandera.tensordict as pa
+
+schema = pa.TensorDictSchema(
+    keys={
+        "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)),
+        "action": pa.Tensor(dtype=torch.int64, shape=(None,)),
+        "reward": pa.Tensor(dtype=torch.float32, shape=(None,), checks=Check.greater_than(0.0)),
+    },
+    batch_size=(32,),
+)
+
+# Save schema to YAML
+schema_yaml = pa.to_yaml(schema)
+print("YAML output preview:", schema_yaml[:150] + "...")
+
+# Load schema from YAML
+loaded_schema = pa.from_yaml(schema_yaml)
+assert loaded_schema.batch_size == schema.batch_size
+```
+
+### JSON Format
+
+```{code-cell} python
+import json
+import torch
+from pandera import Check
+import pandera.tensordict as pa
+
+schema = pa.TensorDictSchema(
+    keys={
+        "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)),
+    },
+    batch_size=(32,),
+)
+
+# Serialize to JSON string
+json_str = pa.to_json(schema)
+print("JSON output preview:", json_str[:150] + "...")
+
+# Deserialize from JSON string
+loaded_schema = pa.from_json(json_str)
+```
+
+## File I/O
+
+Save schemas to and load from files:
+
+```{code-cell} python
+import tempfile
+from pathlib import Path
+import torch
+import pandera.tensordict as pa
+
+schema = pa.TensorDictSchema(
+    keys={
+        "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)),
+    },
+    batch_size=(32,),
+)
+
+# Save to file
+with tempfile.TemporaryDirectory() as tmpdir:
+    schema_path = Path(tmpdir) / "schema.yaml"
+    pa.to_yaml(schema, schema_path)
+<<<<<<< HEAD
+
+=======
+
+>>>>>>> pr/pytorch-tensordict-phase3-4
+    # Load from file
+    loaded_schema = pa.from_yaml(schema_path)
+    print("Successfully saved and loaded schema")
+```
+
+## TensorDict Saving/Loading
+
+Save TensorDicts with embedded schema metadata for data integrity:
+
+```{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.int64, shape=(None,)),
+    },
+    batch_size=(32,),
+)
+
+td = TensorDict({
+    "observation": torch.randn(32, 10),
+    "action": torch.randint(0, 5, (32,)),
+}, batch_size=[32])
+
+# Save with schema metadata
+import tempfile
+with tempfile.TemporaryDirectory() as tmpdir:
+    save_path = f"{tmpdir}/rl_batch.pt"
+    pa.save(schema, td, save_path)
+<<<<<<< HEAD
+
+=======
+
+>>>>>>> pr/pytorch-tensordict-phase3-4
+    # Load and validate automatically
+    loaded_td = pa.load(save_path)
+    print(f"Loaded batch size: {loaded_td.batch_size}")
+```
+
+## Use Cases
+
+### 1. Data Pipeline Validation
+
+Save trained model inputs/outputs with schema:
+
+```{code-cell} python
+import torch
+from tensordict import TensorDict
+import pandera.tensordict as pa
+import tempfile
+
+# At training time
+training_data = TensorDict({
+    "input": torch.randn(100, 64),
+    "target": torch.randint(0, 10, (100,)),
+}, batch_size=[100])
+
+schema = pa.TensorDictSchema(
+    keys={
+        "input": pa.Tensor(dtype=torch.float32, shape=(None, 64)),
+        "target": pa.Tensor(dtype=torch.int64, shape=(None,)),
+    },
+    batch_size=(100,),
+)
+
+with tempfile.TemporaryDirectory() as tmpdir:
+    pa.save(schema, training_data, f"{tmpdir}/training.pt")
+<<<<<<< HEAD
+
+=======
+
+>>>>>>> pr/pytorch-tensordict-phase3-4
+    # Later: validate before inference
+    loaded = pa.load(f"{tmpdir}/training.pt")
+    print("Successfully validated and loaded training data")
+```
+
+### 2. Configuration Management
+
+Define schemas in YAML for version control and collaboration:
+
+```{code-cell} python
+import tempfile
+from pathlib import Path
+import pandera.tensordict as pa
+import torch
+
+# Define schema programmatically
+schema = pa.TensorDictSchema(
+    keys={
+        "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)),
+        "action": pa.Tensor(dtype=torch.int64, shape=(None,)),
+    },
+    batch_size=(32,),
+)
+
+with tempfile.TemporaryDirectory() as tmpdir:
+    # Save to config file
+    config_path = Path(tmpdir) / "rl_schema.yaml"
+    pa.to_yaml(schema, config_path)
+<<<<<<< HEAD
+
+=======
+
+>>>>>>> pr/pytorch-tensordict-phase3-4
+    # Load from config file (e.g., in a different process)
+    loaded_schema = pa.from_yaml(config_path)
+    print("Successfully loaded schema from config")
+```
+
+### 3. Distributed Training
+
+Share schemas across workers via serialization:
+
+```{code-cell} python
+import torch
+import pandera.tensordict as pa
+
+schema = pa.TensorDictSchema(
+    keys={
+        "observation": pa.Tensor(dtype=torch.float32, shape=(None, 10)),
+    },
+    batch_size=(32,),
+)
+
+# Worker 1: Serialize and send
+serialized = pa.to_json(schema)
+print("Serialized schema (JSON):", serialized[:150] + "...")
+
+# Worker 2: Deserialize and use (simulated here)
+loaded_schema = pa.from_json(serialized)
+print(f"Loaded batch_size: {loaded_schema.batch_size}")
+```
+
+## See also
+
+- {ref}`pytorch-tensordict-inference` — infer schemas from data
+- {ref}`pytorch-tensordict-schema` — manual schema definition

docs/source/pytorch_guide/tensordict_model.md

@@ -90,14 +90,45 @@ class RLWithBatchSize(pa.TensorDictModel):
 Use `lazy=True` to collect all validation errors:
 
 ```{code-cell} python
+from tensordict import TensorDict
+
+# Create invalid data with wrong dtypes
+td_wrong_dtype = TensorDict(
+    {"observation": torch.randn(32, 10).to(torch.float64), "action": torch.randint(0, 4, (32,)), "reward": torch.randn(32)},
+    batch_size=[32],
+)
+
 try:
-    RL.validate(td, lazy=True)
+    RL.validate(td_wrong_dtype, 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}")
 ```
 
+## Type Coercion
+
+Model schemas support dtype coercion with the `coerce=True` option:
+
+```{code-cell} python
+class RLWithCoercion(pa.TensorDictModel):
+    observation: torch.float32 = pa.Field(shape=(None, 10))
+
+    class Config:
+        batch_size = (32,)
+        coerce = True
+
+# Input with wrong dtype
+td = TensorDict(
+    {"observation": torch.randn(32, 10).to(torch.float64)},
+    batch_size=[32],
+)
+
+# Dtype is automatically coerced during validation
+validated = RLWithCoercion.validate(td)
+assert validated["observation"].dtype == torch.float32
+```
+
 ## Model inheritance
 
 Models can be inherited to create more specific schemas: