json2vec 0.4.8__tar.gz → 0.4.9__tar.gz

{json2vec-0.4.8/src/json2vec.egg-info → json2vec-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: json2vec
-Version: 0.4.8
+Version: 0.4.9
 Summary: Schema-first PyTorch models for hierarchical / nested / sequence data structures
 License-Expression: Apache-2.0
 Requires-Python: >=3.12

{json2vec-0.4.8 → json2vec-0.4.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "json2vec"
-version = "0.4.8"
+version = "0.4.9"
 description = "Schema-first PyTorch models for hierarchical / nested / sequence data structures"
 readme = "README.md"
 license = "Apache-2.0"

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/__init__.py

@@ -8,15 +8,16 @@ mutation predicates, and the `@preprocess` decorator.
 
 from typing import TYPE_CHECKING, Any
 
+from json2vec import helpers as helpers
+from json2vec.architecture.checkpoint import RollbackCheckpoint
+from json2vec.architecture.mutations import MutationLockCallback, RuntimePlacementCallback
 from json2vec.architecture.root import (
     Model,
-    MutationLockCallback,
     OptimizerConfig,
-    RollbackCheckpoint,
-    RuntimePlacementCallback,
     SchedulerConfig,
 )
 from json2vec.data.datasets import CustomDataModule, PolarsDataModule, StreamingDataModule
+from json2vec.data.processing import MASK_LITERAL, MaskLiteral
 from json2vec.inference.callback import Postprocessor, Writer
 from json2vec.preprocessors import PREPROCESSORS, Preprocessor, PreprocessorMode, preprocess
 from json2vec.structs.enums import (
@@ -38,7 +39,7 @@ from json2vec.structs.experiment import (
     predicate,
     where,
 )
-from json2vec.structs.structure import Array
+from json2vec.structs.structure import Array, Mask
 from json2vec.structs.tree import Address, Leaf
 from json2vec.tensorfields import TENSORFIELDS, DecoderBase, EmbedderBase, Plugin, RequestBase, TensorFieldBase
 from json2vec.tensorfields.extensions.category import Request as Category
@@ -105,11 +106,15 @@ __all__ = [
     "Deployment",
     "EmbedderBase",
     "Entity",
+    "helpers",
     "Hyperparameters",
     "Input",
     "JSONBackend",
     "Leaf",
     "Metric",
+    "MASK_LITERAL",
+    "Mask",
+    "MaskLiteral",
     "Model",
     "ModelSource",
     "MutationLockCallback",

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/architecture/checkpoint.py

@@ -5,7 +5,9 @@ from __future__ import annotations
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import lightning.pytorch as lit
 import torch
+from lightning.pytorch.callbacks import ModelCheckpoint
 from loguru import logger
 
 from json2vec.architecture.graph import ModelGraph
@@ -15,6 +17,46 @@ if TYPE_CHECKING:
     from json2vec.architecture.root import Model
 
 
+class RollbackCheckpoint(ModelCheckpoint):
+    """Checkpoint the best model during fit and restore it into the module at fit end."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        if self.save_weights_only:
+            raise ValueError("RollbackCheckpoint requires full checkpoints; set save_weights_only=False")
+        if self.save_top_k == 0:
+            raise ValueError("RollbackCheckpoint requires at least one saved checkpoint; set save_top_k != 0")
+
+    def on_fit_end(self, trainer: lit.Trainer, pl_module: lit.LightningModule) -> None:
+        from json2vec.architecture.root import Model
+
+        super().on_fit_end(trainer=trainer, pl_module=pl_module)
+        if not isinstance(pl_module, Model):
+            raise TypeError("RollbackCheckpoint can only restore json2vec Model instances")
+
+        best_model_path = self.best_model_path
+        if not best_model_path:
+            raise RuntimeError("RollbackCheckpoint did not find a best checkpoint to restore")
+
+        strategy = getattr(trainer, "strategy", None)
+        if strategy is not None:
+            strategy.barrier("rollback_checkpoint_load")
+            checkpoint = strategy.checkpoint_io.load_checkpoint(
+                best_model_path,
+                map_location=pl_module.device,
+                weights_only=False,
+            )
+        else:
+            checkpoint = torch.load(best_model_path, weights_only=False, map_location=pl_module.device)
+
+        pl_module.restore_checkpoint_state(checkpoint)
+        logger.bind(
+            component="checkpoint",
+            checkpoint=best_model_path,
+            score=self.best_model_score,
+        ).info("rolled back Model to best checkpoint")
+
+
 class CheckpointState:
     """Save, load, and restore model state without owning the public facade."""
 

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/architecture/mutations.py

@@ -4,12 +4,17 @@ from __future__ import annotations
 
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
+from functools import partialmethod, wraps
 from typing import TYPE_CHECKING, Any
 
+import lightning.pytorch as lit
 import pydantic
+import torch
+from lightning.pytorch import Callback
 from loguru import logger
 
 from json2vec.architecture.graph import ModelGraph
+from json2vec.structs.enums import Strata
 from json2vec.structs.experiment import NodeAttribute, NodePredicate, SchemaField
 from json2vec.structs.structure import Array
 from json2vec.structs.tree import Leaf, Node
@@ -20,6 +25,73 @@ if TYPE_CHECKING:
 _MISSING = object()
 
 
+def immutable(name: str | Strata) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    def decorator(method: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(method)
+        def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
+            locks = self.locks
+            locks[name] += 1
+            try:
+                return method(self, *args, **kwargs)
+            finally:
+                if locks[name] <= 1:
+                    locks.pop(name, None)
+                else:
+                    locks[name] -= 1
+
+        return wrapped
+
+    return decorator
+
+
+class MutationLockCallback(Callback):
+    """Prevent runtime schema mutations while Lightning owns an active loop."""
+
+    locks: tuple[Strata, ...] = (Strata.train, Strata.validate, Strata.test, Strata.predict)
+
+    def _on_loop_start(self, trainer: lit.Trainer, pl_module: "Model", strata: Strata) -> None:
+        pl_module.locks[strata] += 1
+
+    def _on_loop_end(self, trainer: lit.Trainer, pl_module: "Model", strata: Strata) -> None:
+        locks = pl_module.locks
+        if locks[strata] <= 1:
+            locks.pop(strata, None)
+        else:
+            locks[strata] -= 1
+
+    def on_exception(
+        self,
+        trainer: lit.Trainer,
+        pl_module: "Model",
+        exception: BaseException,
+    ) -> None:  # ty:ignore[invalid-method-override]
+        for lock in self.locks:
+            pl_module.locks.pop(lock, None)
+
+    on_train_start = partialmethod(_on_loop_start, strata=Strata.train)
+    on_train_end = partialmethod(_on_loop_end, strata=Strata.train)
+    on_validation_start = partialmethod(_on_loop_start, strata=Strata.validate)
+    on_validation_end = partialmethod(_on_loop_end, strata=Strata.validate)
+    on_test_start = partialmethod(_on_loop_start, strata=Strata.test)
+    on_test_end = partialmethod(_on_loop_end, strata=Strata.test)
+    on_predict_start = partialmethod(_on_loop_start, strata=Strata.predict)
+    on_predict_end = partialmethod(_on_loop_end, strata=Strata.predict)
+
+
+class RuntimePlacementCallback(Callback):
+    """Move late-created modules onto the Lightning module's active device."""
+
+    def _on_loop_start(self, trainer: lit.Trainer, pl_module: lit.LightningModule, strata: Strata) -> None:
+        device = getattr(pl_module, "device", None)
+        if isinstance(device, torch.device):
+            pl_module.to(device=device)
+
+    on_train_start = partialmethod(_on_loop_start, strata=Strata.train)
+    on_validation_start = partialmethod(_on_loop_start, strata=Strata.validate)
+    on_test_start = partialmethod(_on_loop_start, strata=Strata.test)
+    on_predict_start = partialmethod(_on_loop_start, strata=Strata.predict)
+
+
 class AttributeChange(pydantic.BaseModel):
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
@@ -40,6 +112,12 @@ class SchemaEditor:
     def __init__(self, module: "Model") -> None:
         self.module = module
 
+    def _assert_mutation_allowed(self, action: str) -> None:
+        active = tuple(name for name, count in self.module.locks.items() if count > 0)
+        if active:
+            labels = ", ".join(active)
+            raise RuntimeError(f"model.{action}(...) cannot run while the model is in an active loop: {labels}")
+
     def select(
         self,
         *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
@@ -62,7 +140,7 @@ class SchemaEditor:
         use_cache: bool = False,
         **values: Any,
     ) -> None:
-        self.module._assert_mutation_allowed("update")
+        self._assert_mutation_allowed("update")
         values = self.module.hyperparameters.update_values(values)
         changes = self._attribute_changes(
             values=values,
@@ -90,7 +168,7 @@ class SchemaEditor:
         include_root: bool = True,
         use_cache: bool = True,
     ) -> None:
-        self.module._assert_mutation_allowed("extend")
+        self._assert_mutation_allowed("extend")
         parent, field_count = self._extend_target(*args, include_root=include_root, use_cache=use_cache)
         self.module.hyperparameters.extend(*args, include_root=include_root, use_cache=use_cache)
         ModelGraph.rebuild(self.module)
@@ -109,7 +187,7 @@ class SchemaEditor:
         include_root: bool = False,
         use_cache: bool = True,
     ) -> None:
-        self.module._assert_mutation_allowed("delete")
+        self._assert_mutation_allowed("delete")
         roots = self._delete_roots(*predicates, include_root=include_root, use_cache=use_cache)
         self.module.hyperparameters.delete(*predicates, include_root=include_root, use_cache=use_cache)
         ModelGraph.rebuild(self.module)
@@ -129,7 +207,7 @@ class SchemaEditor:
         use_cache: bool = True,
         descendants: bool = False,
     ) -> None:
-        self.module._assert_mutation_allowed("reset")
+        self._assert_mutation_allowed("reset")
         selected = self.module.hyperparameters.select(
             *predicates,
             include_root=include_root,
@@ -160,7 +238,7 @@ class SchemaEditor:
         use_cache: bool = False,
         **values: Any,
     ) -> Iterator[None]:
-        self.module._assert_mutation_allowed("override")
+        self._assert_mutation_allowed("override")
         values = self.module.hyperparameters.update_values(values)
         changes = self._attribute_changes(
             values=values,

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/architecture/root.py

@@ -3,7 +3,7 @@
 from collections import Counter
 from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager
-from functools import partialmethod, wraps
+from functools import partialmethod
 from pathlib import Path
 from typing import Any, Self, cast
 
@@ -11,15 +11,19 @@ import lightning.pytorch as lit
 import torch
 from beartype import beartype
 from lightning.pytorch import Callback
-from lightning.pytorch.callbacks import ModelCheckpoint
 from loguru import logger
 from rich.text import Text
 from tensordict import TensorDict
 
-from json2vec.architecture.checkpoint import CheckpointState
+from json2vec.architecture.checkpoint import CheckpointState, RollbackCheckpoint
 from json2vec.architecture.contracts import ContractScheduler
 from json2vec.architecture.graph import ModelGraph
-from json2vec.architecture.mutations import SchemaEditor
+from json2vec.architecture.mutations import (
+    MutationLockCallback,
+    RuntimePlacementCallback,
+    SchemaEditor,
+    immutable,
+)
 from json2vec.architecture.runtime import ModelRuntime, Postprocessor, Preprocessor, step
 from json2vec.data.datasets.base import EncodedBatch, EncodedInput
 from json2vec.logging.throughput import ThroughputLogger
@@ -37,105 +41,12 @@ from json2vec.tensorfields.base import TENSORFIELDS, Plugin, TensorFieldBase
 OptimizerConfig = torch.optim.Optimizer | Callable[["Model"], torch.optim.Optimizer]
 SchedulerConfig = Any | Callable[["Model", torch.optim.Optimizer], Any]
 
-
-def immutable(name: str | Strata) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
-    def decorator(method: Callable[..., Any]) -> Callable[..., Any]:
-        @wraps(method)
-        def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
-            locks = self.locks
-            locks[name] += 1
-            try:
-                return method(self, *args, **kwargs)
-            finally:
-                if locks[name] <= 1:
-                    locks.pop(name, None)
-                else:
-                    locks[name] -= 1
-
-        return wrapped
-
-    return decorator
-
-
-class MutationLockCallback(Callback):
-    """Prevent runtime schema mutations while Lightning owns an active loop."""
-
-    locks: tuple[Strata, ...] = (Strata.train, Strata.validate, Strata.test, Strata.predict)
-
-    def _on_loop_start(self, trainer: lit.Trainer, pl_module: "Model", strata: Strata) -> None:
-        pl_module.locks[strata] += 1
-
-    def _on_loop_end(self, trainer: lit.Trainer, pl_module: "Model", strata: Strata) -> None:
-        locks = pl_module.locks
-        if locks[strata] <= 1:
-            locks.pop(strata, None)
-        else:
-            locks[strata] -= 1
-
-    def on_exception(self, trainer: lit.Trainer, pl_module: "Model", exception: BaseException) -> None:  # ty:ignore[invalid-method-override]
-        for lock in self.locks:
-            pl_module.locks.pop(lock, None)
-
-    on_train_start = partialmethod(_on_loop_start, strata=Strata.train)
-    on_train_end = partialmethod(_on_loop_end, strata=Strata.train)
-    on_validation_start = partialmethod(_on_loop_start, strata=Strata.validate)
-    on_validation_end = partialmethod(_on_loop_end, strata=Strata.validate)
-    on_test_start = partialmethod(_on_loop_start, strata=Strata.test)
-    on_test_end = partialmethod(_on_loop_end, strata=Strata.test)
-    on_predict_start = partialmethod(_on_loop_start, strata=Strata.predict)
-    on_predict_end = partialmethod(_on_loop_end, strata=Strata.predict)
-
-
-class RuntimePlacementCallback(Callback):
-    """Move late-created modules onto the Lightning module's active device."""
-
-    def _on_loop_start(self, trainer: lit.Trainer, pl_module: lit.LightningModule, strata: Strata) -> None:
-        device = getattr(pl_module, "device", None)
-        if isinstance(device, torch.device):
-            pl_module.to(device=device)
-
-    on_train_start = partialmethod(_on_loop_start, strata=Strata.train)
-    on_validation_start = partialmethod(_on_loop_start, strata=Strata.validate)
-    on_test_start = partialmethod(_on_loop_start, strata=Strata.test)
-    on_predict_start = partialmethod(_on_loop_start, strata=Strata.predict)
-
-
-class RollbackCheckpoint(ModelCheckpoint):
-    """Checkpoint the best model during fit and restore it into the module at fit end."""
-
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        if self.save_weights_only:
-            raise ValueError("RollbackCheckpoint requires full checkpoints; set save_weights_only=False")
-        if self.save_top_k == 0:
-            raise ValueError("RollbackCheckpoint requires at least one saved checkpoint; set save_top_k != 0")
-
-    def on_fit_end(self, trainer: lit.Trainer, pl_module: lit.LightningModule) -> None:
-        super().on_fit_end(trainer=trainer, pl_module=pl_module)
-        if not isinstance(pl_module, Model):
-            raise TypeError("RollbackCheckpoint can only restore json2vec Model instances")
-
-        best_model_path = self.best_model_path
-        if not best_model_path:
-            raise RuntimeError("RollbackCheckpoint did not find a best checkpoint to restore")
-
-        strategy = getattr(trainer, "strategy", None)
-        if strategy is not None:
-            strategy.barrier("rollback_checkpoint_load")
-            checkpoint = strategy.checkpoint_io.load_checkpoint(
-                best_model_path,
-                map_location=pl_module.device,
-                weights_only=False,
-            )
-        else:
-            checkpoint = torch.load(best_model_path, weights_only=False, map_location=pl_module.device)
-
-        pl_module.restore_checkpoint_state(checkpoint)
-        logger.bind(
-            component="checkpoint",
-            checkpoint=best_model_path,
-            score=self.best_model_score,
-        ).info("rolled back Model to best checkpoint")
+__all__ = [
+    "Model",
+    "MutationLockCallback",
+    "RollbackCheckpoint",
+    "RuntimePlacementCallback",
+]
 
 
 class Model(lit.LightningModule, Renderable):
@@ -406,12 +317,6 @@ class Model(lit.LightningModule, Renderable):
         ):
             yield
 
-    def _assert_mutation_allowed(self, action: str) -> None:
-        active = tuple(name for name, count in self.locks.items() if count > 0)
-        if active:
-            labels = ", ".join(active)
-            raise RuntimeError(f"model.{action}(...) cannot run while the model is in an active loop: {labels}")
-
     def configure_callbacks(self) -> list[Callback]:
         callbacks: list[Callback] = []
         factories: set[Any] = set()
@@ -538,6 +443,7 @@ class Model(lit.LightningModule, Renderable):
         batch: EncodedBatch | list[dict[str, Any]],
         preprocess: Preprocessor | None = None,
         strata: Strata | str = Strata.predict,
+        mask: bool = True,
     ) -> EncodedInput:
         """Return encoded tensorfield inputs for raw or processed observations."""
         return ModelRuntime.encode(
@@ -545,6 +451,7 @@ class Model(lit.LightningModule, Renderable):
             batch=batch,
             preprocess=preprocess,
             strata=strata,
+            mask=mask,
         )
 
     @immutable("inference")

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/architecture/runtime.py

@@ -14,8 +14,9 @@ from json2vec.architecture.contracts import sanitize
 from json2vec.architecture.encoder import ArrayEncoder
 from json2vec.architecture.node import NodeModule
 from json2vec.data.datasets.base import EncodedBatch, EncodedInput
-from json2vec.data.iterables import encode
-from json2vec.structs.enums import Metric, Strata, TensorKey
+from json2vec.data.iterables import encode as encode_batch
+from json2vec.data.iterables import mask as apply_mask
+from json2vec.structs.enums import Metric, Strata, TensorKey, Tokens
 from json2vec.structs.packages import Parcel, Prediction
 from json2vec.structs.tree import Address
 from json2vec.tensorfields.base import (
@@ -99,8 +100,10 @@ class ModelRuntime:
                     )
 
         for address in module.hyperparameters.active_requests.keys():
+            has_masked_input = inputs[address].state.eq(Tokens.masked.value).any()
             if (
                 torch.any(inputs[address].trainable)
+                or (strata == Strata.predict and has_masked_input)
                 or (address in module.hyperparameters.target)
                 or (address in module.hyperparameters.embed)
             ):
@@ -193,6 +196,7 @@ class ModelRuntime:
         batch: EncodedBatch | list[dict[str, Any]],
         preprocess: Preprocessor | None = None,
         strata: Strata | str = Strata.predict,
+        mask: bool = True,
     ) -> EncodedInput:
         strata = Strata.normalize(strata)
 
@@ -209,12 +213,17 @@ class ModelRuntime:
         elif batch and isinstance(batch[0], dict):
             batch = [[request] for request in cast(list[dict[str, Any]], batch)]
 
-        return encode(
+        inputs = encode_batch(
             batch=cast(EncodedBatch, batch),
             hyperparameters=module.hyperparameters,
             strata=strata,
             interprocess_encoding_context=module.interprocess_encoding_context,
+            defer_target_masking=True,
         )
+        if mask:
+            return next(apply_mask([inputs], module.hyperparameters, strata=strata))
+
+        return inputs
 
     @staticmethod
     def predict(

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/data/datasets/custom.py

@@ -31,7 +31,6 @@ from json2vec.data.iterables import (
     process,
     sample,
     shuffle,
-    target,
     transform,
 )
 from json2vec.data.processing import Pipeline
@@ -130,7 +129,6 @@ class CustomBatchDataset(IterableDataset):
             | batch
             | transform
             | mask
-            | target
         )
 
 

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/data/datasets/polars.py

@@ -35,7 +35,6 @@ from json2vec.data.iterables import (
     process,
     sample,
     shuffle,
-    target,
     transform,
 )
 from json2vec.data.processing import Pipeline
@@ -183,7 +182,6 @@ class PolarsBatchDataset(IterableDataset):
             | batch
             | transform
             | mask
-            | target
         )
 
 

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/data/datasets/streaming.py

@@ -40,7 +40,6 @@ from json2vec.data.iterables import (
     process,
     sample,
     shuffle,
-    target,
     transform,
 )
 from json2vec.data.processing import Pipeline
@@ -323,7 +322,6 @@ class BatchDataset(IterableDataset):
             | batch
             | transform
             | mask
-            | target
         )
 
 

{json2vec-0.4.8 → json2vec-0.4.9}/src/json2vec/data/iterables.py

@@ -23,6 +23,7 @@ from json2vec.data.datasets.base import (
     ProcessedObservation,
     RawObservation,
 )
+from json2vec.data.processing import MASK_LITERAL, contains_mask_literal
 from json2vec.preprocessors.base import PREPROCESSORS, Preprocessor, PreprocessorMode
 from json2vec.structs.enums import Strata, TensorKey
 from json2vec.structs.experiment import Hyperparameters
@@ -224,10 +225,14 @@ def encode(
     strata: Strata,
     interprocess_encoding_context: InterprocessEncodingContext,
     jmespath_resolution_monitor: JMESPathResolutionMonitor | None = None,
+    defer_target_masking: bool = False,
 ) -> EncodedInput:
     out: dict[Address, TensorFieldBase] = {}
     target_addresses = set(hyperparameters.target)
 
+    if strata != Strata.predict and contains_mask_literal(batch):
+        raise ValueError(f"{MASK_LITERAL!r} is only valid during predict strata")
+
     for address, request in hyperparameters.active_requests.items():
         TensorField = cast(type[TensorFieldBase], getattr(TENSORFIELDS[request.type], "TensorField"))
 
@@ -262,8 +267,8 @@ def encode(
 
         out[address] = TensorField.new(**kwargs)
 
-        if address in target_addresses:
-            out[address].target(p_prune=1.0)
+        if not defer_target_masking and strata != Strata.predict and address in target_addresses:
+            out[address].mask(p_prune=1.0)
 
     inputs = cast(EncodedInput, TensorDict(source=cast(Any, out)))
 
@@ -288,21 +293,69 @@ def transform(
             strata=strata,
             interprocess_encoding_context=interprocess_encoding_context,
             jmespath_resolution_monitor=jmespath_resolution_monitor,
+            defer_target_masking=True,
         )
 
 
+def _apply_mask_policy(
+    field: TensorFieldBase,
+    *,
+    p_mask: float,
+    p_prune: float,
+    array_masks: tuple[Any, ...],
+    address: Address,
+    hyperparameters: Hyperparameters,
+) -> None:
+    parameters = inspect.signature(field.mask).parameters
+    supports_policy_kwargs = any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values())
+    supports_policy_kwargs |= any(name in parameters for name in ("p_prune", "array_masks", "hyperparameters"))
+
+    if supports_policy_kwargs:
+        field.mask(
+            p_mask=p_mask,
+            p_prune=p_prune,
+            array_masks=array_masks,
+            address=address,
+            hyperparameters=hyperparameters,
+        )
+        return
+
+    if array_masks:
+        raise TypeError(f"tensorfield at '{address}' must accept mask(..., array_masks=...) to use Array masks")
+
+    if p_mask > 0.0:
+        field.mask(p_mask=p_mask)
+
+    if p_prune > 0.0:
+        field.target(p_prune=p_prune)
+
+
 @beartype
 def mask(
     pipe: Iterable[EncodedInput],
     hyperparameters: Hyperparameters,
+    strata: Strata = Strata.train,
 ) -> Iterator[EncodedInput]:
     for item in pipe:
+        if strata == Strata.predict:
+            yield item
+            continue
+
         for address, request in hyperparameters.active_requests.items():
             p_mask = float(request.p_mask or 0.0)
-            if p_mask <= 0.0:
+            p_prune = float(request.p_prune or 0.0)
+            array_masks = hyperparameters.array_masks_for(address)
+            if p_mask <= 0.0 and p_prune <= 0.0 and not array_masks:
                 continue
 
-            item[address].mask(p_mask=p_mask)
+            _apply_mask_policy(
+                item[address],
+                p_mask=p_mask,
+                p_prune=p_prune,
+                array_masks=array_masks,
+                address=address,
+                hyperparameters=hyperparameters,
+            )
 
         yield item