From e13cbe5827dac20f7307049871e3f50f4b5411e7 Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Thu, 15 Oct 2026 13:35:12 -0700
Subject: [PATCH 1/7] Add torch and lsdb dependencies and device helpers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 pyproject.toml               |  2 +
 src/lsdb_ml/device.py        | 77 ++++++++++++++++++++++++++++++++++++
 tests/lsdb_ml/test_device.py | 26 ++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 src/lsdb_ml/device.py
 create mode 100644 tests/lsdb_ml/test_device.py

diff --git a/pyproject.toml b/pyproject.toml
index 799fdb2..8f999d0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,6 +17,8 @@ classifiers = [
 dynamic = ["version"]
 requires-python = ">=3.11"
 dependencies = [
+    "lsdb>=0.11.0", # Brings nested-pandas, pyarrow and dask[complete] (incl. distributed)
+    "torch>=2.4",
 ]
 
 [project.urls]
diff --git a/src/lsdb_ml/device.py b/src/lsdb_ml/device.py
new file mode 100644
index 0000000..3dc0c55
--- /dev/null
+++ b/src/lsdb_ml/device.py
@@ -0,0 +1,77 @@
+"""Device resolution and batch movement helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+
+DeviceLike = str | torch.device | None
+
+
+def resolve_device(device: DeviceLike = "cpu") -> torch.device:
+    """Turn a device request into a concrete ``torch.device``.
+
+    Parameters
+    ----------
+    device : str, torch.device or None, default "cpu"
+        ``"cpu"`` (or ``None``) for the CPU, ``"auto"`` to pick CUDA if available, then
+        MPS, then CPU. Any other value (e.g. ``"cuda"``, ``"cuda:1"``, ``"mps"``, or a
+        ``torch.device``) is passed to ``torch.device`` and checked for availability.
+
+    Returns
+    -------
+    torch.device
+        The resolved device.
+
+    Raises
+    ------
+    RuntimeError
+        If an explicitly requested accelerator is not available.
+
+    Examples
+    --------
+    >>> resolve_device("cpu")
+    device(type='cpu')
+    >>> resolve_device(None)
+    device(type='cpu')
+    """
+    if device is None:
+        return torch.device("cpu")
+    if isinstance(device, str) and device == "auto":
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+
+    resolved = torch.device(device)
+    if resolved.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError(f"Device {resolved} was requested, but CUDA is not available.")
+    if resolved.type == "mps" and not torch.backends.mps.is_available():
+        raise RuntimeError(f"Device {resolved} was requested, but MPS is not available.")
+    return resolved
+
+
+def move_batch(batch: Mapping[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]:
+    """Move every tensor in a batch dict to ``device``.
+
+    Non-tensor values are passed through unchanged.
+
+    Parameters
+    ----------
+    batch : Mapping[str, torch.Tensor]
+        A collated batch, as produced by :meth:`lsdb_ml.ResolvedSpec.collate`.
+    device : torch.device
+        The target device.
+
+    Returns
+    -------
+    dict[str, torch.Tensor]
+        A new dict with the tensors on ``device``.
+    """
+    non_blocking = device.type == "cuda"
+    return {
+        key: value.to(device, non_blocking=non_blocking) if isinstance(value, torch.Tensor) else value
+        for key, value in batch.items()
+    }
diff --git a/tests/lsdb_ml/test_device.py b/tests/lsdb_ml/test_device.py
new file mode 100644
index 0000000..cb90019
--- /dev/null
+++ b/tests/lsdb_ml/test_device.py
@@ -0,0 +1,26 @@
+import pytest
+import torch
+
+from lsdb_ml import move_batch, resolve_device
+
+
+def test_resolve_device():
+    """Known device requests resolve to torch devices."""
+    assert resolve_device("cpu") == torch.device("cpu")
+    assert resolve_device(torch.device("cpu")) == torch.device("cpu")
+    assert resolve_device("auto").type in {"cpu", "cuda", "mps"}
+
+
+@pytest.mark.skipif(torch.cuda.is_available(), reason="checks the error when CUDA is missing")
+def test_resolve_missing_cuda():
+    """Requesting CUDA without CUDA fails clearly."""
+    with pytest.raises(RuntimeError, match="CUDA is not available"):
+        resolve_device("cuda")
+
+
+def test_move_batch():
+    """Tensors move to the device and other values pass through."""
+    batch = {"x": torch.zeros(2), "meta": "keep"}
+    moved = move_batch(batch, torch.device("cpu"))
+    assert moved["x"].device.type == "cpu"
+    assert moved["meta"] == "keep"

From 02c20fef1465133e08007bc465920dd1ba18b926 Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Thu, 15 Oct 2026 13:35:12 -0700
Subject: [PATCH 2/7] Add ColumnSpec: meta validation and padded nested
 collation

ColumnSpec maps flat and nested numeric catalog columns to tensors.
resolve() validates against catalog meta (nested sub-column dtypes come
from NestedDtype.column_dtypes) without computing data, and reports
every problem at once. Nested columns are zero-padded to the batch
maximum with a mask and lengths, vectorized via nest.len() and
nest.to_flat(), with optional max_length truncation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/lsdb_ml/collate.py        | 157 +++++++++++++
 src/lsdb_ml/spec.py           | 423 ++++++++++++++++++++++++++++++++++
 tests/lsdb_ml/conftest.py     |  57 +++++
 tests/lsdb_ml/test_collate.py | 136 +++++++++++
 tests/lsdb_ml/test_spec.py    | 109 +++++++++
 5 files changed, 882 insertions(+)
 create mode 100644 src/lsdb_ml/collate.py
 create mode 100644 src/lsdb_ml/spec.py
 create mode 100644 tests/lsdb_ml/test_collate.py
 create mode 100644 tests/lsdb_ml/test_spec.py

diff --git a/src/lsdb_ml/collate.py b/src/lsdb_ml/collate.py
new file mode 100644
index 0000000..2de8aea
--- /dev/null
+++ b/src/lsdb_ml/collate.py
@@ -0,0 +1,157 @@
+"""Vectorized conversion of (nested) dataframe columns into padded numpy arrays.
+
+These helpers do the numeric work behind :class:`lsdb_ml.Flat` and :class:`lsdb_ml.Nested`.
+Nothing here loops over rows in Python.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+import torch
+
+TruncatePolicy = Literal["raise", "head", "tail"]
+
+
+def numpy_dtype(dtype: torch.dtype) -> np.dtype:
+    """Return the numpy dtype that corresponds to a torch dtype.
+
+    Parameters
+    ----------
+    dtype : torch.dtype
+        A torch dtype with a numpy equivalent.
+
+    Returns
+    -------
+    np.dtype
+        The matching numpy dtype.
+
+    Raises
+    ------
+    TypeError
+        If the torch dtype has no numpy equivalent (e.g. ``torch.bfloat16``).
+
+    Examples
+    --------
+    >>> numpy_dtype(torch.float32)
+    dtype('float32')
+    """
+    try:
+        return torch.empty(0, dtype=dtype).numpy().dtype
+    except TypeError as err:
+        raise TypeError(f"torch dtype {dtype} has no numpy equivalent and cannot be collated.") from err
+
+
+def frame_to_array(frame: pd.DataFrame, dtype: np.dtype, what: str) -> np.ndarray:
+    """Convert the columns of ``frame`` into a writable ``(N, F)`` array.
+
+    Parameters
+    ----------
+    frame : pd.DataFrame
+        The columns to convert; may be arrow- or numpy-backed.
+    dtype : np.dtype
+        Output dtype.
+    what : str
+        Description of the data used in error messages.
+
+    Returns
+    -------
+    np.ndarray
+        A C-contiguous, writable array of shape ``(len(frame), len(frame.columns))``.
+
+    Raises
+    ------
+    ValueError
+        If any value is missing.
+    """
+    missing = frame.isna().to_numpy()
+    if missing.any():
+        n_missing = int(missing.any(axis=1).sum())
+        raise ValueError(
+            f"{what} contains {n_missing} row(s) with missing values. Drop or fill them before "
+            "collation, e.g. with catalog.dropna() or catalog.map_partitions()."
+        )
+    # np.array (not np.asarray) always copies: arrow-backed columns can convert to read-only
+    # views, which torch.from_numpy would warn about and share.
+    return np.array(frame.to_numpy(dtype=dtype), dtype=dtype, order="C", copy=True)
+
+
+def pad_ragged(
+    values: np.ndarray,
+    lengths: np.ndarray,
+    *,
+    max_length: int | None = None,
+    truncate: TruncatePolicy = "raise",
+    what: str = "nested column",
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Scatter flat ragged values into a zero-padded ``(B, L, F)`` array with a mask.
+
+    Parameters
+    ----------
+    values : np.ndarray
+        Flat values of shape ``(sum(lengths), F)``, rows in order.
+    lengths : np.ndarray
+        Number of values in each of the ``B`` rows.
+    max_length : int or None, default None
+        Longest sequence to keep. Longer rows are handled by ``truncate``.
+    truncate : {"raise", "head", "tail"}, default "raise"
+        For rows over ``max_length``: raise, keep the first values, or keep the last values.
+    what : str, default "nested column"
+        Description of the data used in error messages.
+
+    Returns
+    -------
+    padded : np.ndarray
+        Shape ``(B, L, F)``, where ``L`` is the longest kept row.
+    mask : np.ndarray
+        Boolean, shape ``(B, L)``; True where ``padded`` holds real data.
+    kept_lengths : np.ndarray
+        int64, shape ``(B,)``; number of real values per row after truncation.
+
+    Examples
+    --------
+    >>> padded, mask, kept = pad_ragged(np.arange(5.0)[:, None], np.array([2, 0, 3]))
+    >>> padded[..., 0]
+    array([[0., 1., 0.],
+           [0., 0., 0.],
+           [2., 3., 4.]])
+    >>> mask.astype(int)
+    array([[1, 1, 0],
+           [0, 0, 0],
+           [1, 1, 1]])
+    """
+    lengths = np.asarray(lengths, dtype=np.int64)
+    n_rows = len(lengths)
+    n_features = values.shape[1]
+    if int(lengths.sum()) != len(values):
+        raise ValueError(f"{what}: lengths sum to {int(lengths.sum())} but there are {len(values)} values.")
+
+    starts = np.cumsum(lengths) - lengths
+    kept = lengths
+    source_starts = starts
+    if max_length is not None and n_rows and int(lengths.max()) > max_length:
+        if truncate == "raise":
+            raise ValueError(
+                f"{what} has rows of length up to {int(lengths.max())}, over max_length={max_length}. "
+                "Pass truncate='head' or truncate='tail' to cut them."
+            )
+        kept = np.minimum(lengths, max_length)
+        if truncate == "tail":
+            source_starts = starts + lengths - kept
+        elif truncate != "head":
+            raise ValueError(f"truncate must be 'raise', 'head' or 'tail', got {truncate!r}.")
+
+    width = int(kept.max()) if n_rows else 0
+    padded = np.zeros((n_rows, width, n_features), dtype=values.dtype)
+    mask = np.zeros((n_rows, width), dtype=bool)
+
+    total = int(kept.sum())
+    if total:
+        row = np.repeat(np.arange(n_rows), kept)
+        position = np.arange(total) - np.repeat(np.cumsum(kept) - kept, kept)
+        source = np.repeat(source_starts, kept) + position
+        padded[row, position] = values[source]
+        mask[row, position] = True
+    return padded, mask, kept.astype(np.int64)
diff --git a/src/lsdb_ml/spec.py b/src/lsdb_ml/spec.py
new file mode 100644
index 0000000..f215d3e
--- /dev/null
+++ b/src/lsdb_ml/spec.py
@@ -0,0 +1,423 @@
+"""Column specifications: map flat and nested catalog columns to tensors.
+
+A :class:`ColumnSpec` names the tensors a model needs and where each comes from. It is
+checked against a catalog's ``meta`` (an empty ``NestedFrame`` that carries every dtype,
+including nested sub-column dtypes) before any data is computed, and the resulting
+:class:`ResolvedSpec` collates dataframe rows into a ``dict[str, torch.Tensor]``.
+
+>>> import lsdb_ml
+>>> spec = lsdb_ml.ColumnSpec({"x": lsdb_ml.Flat(["ra", "dec"]), "lc": lsdb_ml.Nested("lc", ["flux"])})
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import torch
+from nested_pandas.series.dtype import NestedDtype
+
+from lsdb_ml.collate import TruncatePolicy, frame_to_array, numpy_dtype, pad_ragged
+
+
+class SpecError(ValueError):
+    """Raised when a column specification does not match a catalog's schema."""
+
+
+@dataclass(frozen=True)
+class TensorInfo:
+    """Description of one tensor a resolved spec produces.
+
+    Attributes
+    ----------
+    dtype : torch.dtype
+        The tensor dtype.
+    shape : tuple
+        The shape; ``None`` marks a dimension that varies by batch (always the first one).
+    """
+
+    dtype: torch.dtype
+    shape: tuple[int | None, ...]
+
+
+def _as_list(columns: str | Sequence[str]) -> list[str]:
+    return [columns] if isinstance(columns, str) else list(columns)
+
+
+def _is_numeric(dtype) -> bool:
+    """True for pandas/numpy/pyarrow integer, float, and bool dtypes."""
+    if isinstance(dtype, pa.DataType):
+        return pa.types.is_integer(dtype) or pa.types.is_floating(dtype) or pa.types.is_boolean(dtype)
+    if isinstance(dtype, pd.ArrowDtype):
+        return _is_numeric(dtype.pyarrow_dtype)
+    if isinstance(dtype, NestedDtype):
+        return False
+    try:
+        return pd.api.types.is_numeric_dtype(dtype) or pd.api.types.is_bool_dtype(dtype)
+    except TypeError:
+        return False
+
+
+class Field(ABC):
+    """Base class for one entry of a :class:`ColumnSpec`.
+
+    Subclasses implement :meth:`resolve`, which validates the field against catalog meta
+    and returns a :class:`ResolvedField` that does the actual collation. New column
+    representations (e.g. ragged tensors) are added as new ``Field`` types.
+    """
+
+    @abstractmethod
+    def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None:
+        """Validate against ``meta`` and build the collator.
+
+        Parameters
+        ----------
+        key : str
+            The output key this field was registered under.
+        meta : pd.DataFrame
+            The catalog meta (an empty frame with full dtypes).
+        problems : list[str]
+            Append a message for every problem found, instead of raising, so that all
+            problems in a spec are reported together.
+
+        Returns
+        -------
+        ResolvedField or None
+            None if any problem was found.
+        """
+
+
+class ResolvedField(ABC):
+    """A field validated against a schema, ready to collate frames."""
+
+    #: Top-level columns this field reads.
+    columns: list[str]
+
+    @abstractmethod
+    def describe(self) -> dict[str, TensorInfo]:
+        """Return the tensors this field produces, keyed by output key."""
+
+    @abstractmethod
+    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
+        """Convert the rows of ``frame`` to tensors."""
+
+
+@dataclass
+class Flat(Field):
+    """One or more flat (non-nested) numeric columns, stacked into a single tensor.
+
+    Parameters
+    ----------
+    columns : str or sequence of str
+        A single column name produces a ``(B,)`` tensor; a list produces ``(B, F)``.
+    dtype : torch.dtype, default torch.float32
+        Output dtype.
+
+    Notes
+    -----
+    Missing values raise an error during collation; drop or fill them in the catalog first.
+    """
+
+    columns: str | Sequence[str]
+    dtype: torch.dtype = torch.float32
+
+    def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None:
+        """Validate the columns against ``meta``; see :meth:`Field.resolve`."""
+        columns = _as_list(self.columns)
+        start = len(problems)
+        if not columns:
+            problems.append(f"'{key}': Flat needs at least one column.")
+        for column in columns:
+            if column not in meta.columns:
+                problems.append(f"'{key}': column '{column}' is not in the catalog.")
+            elif isinstance(meta.dtypes[column], NestedDtype):
+                problems.append(f"'{key}': column '{column}' is nested; use Nested('{column}', [...]).")
+            elif not _is_numeric(meta.dtypes[column]):
+                problems.append(f"'{key}': column '{column}' has non-numeric dtype {meta.dtypes[column]}.")
+        try:
+            np_dtype = numpy_dtype(self.dtype)
+        except TypeError as err:
+            problems.append(f"'{key}': {err}")
+        if len(problems) > start:
+            return None
+        return _ResolvedFlat(key, columns, isinstance(self.columns, str), self, np_dtype)
+
+
+class _ResolvedFlat(ResolvedField):
+    def __init__(self, key: str, columns: list[str], scalar: bool, spec: Flat, np_dtype: np.dtype):
+        self.key = key
+        self.columns = columns
+        self.scalar = scalar
+        self.spec = spec
+        self.np_dtype = np_dtype
+
+    def describe(self) -> dict[str, TensorInfo]:
+        shape = (None,) if self.scalar else (None, len(self.columns))
+        return {self.key: TensorInfo(self.spec.dtype, shape)}
+
+    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
+        array = frame_to_array(frame[self.columns], self.np_dtype, f"'{self.key}' (columns {self.columns})")
+        if self.scalar:
+            array = array[:, 0].copy()
+        return {self.key: torch.from_numpy(array)}
+
+
+@dataclass
+class Nested(Field):
+    """Numeric sub-columns of a nested column, zero-padded to a rectangular tensor with a mask.
+
+    For output key ``k`` this produces:
+
+    - ``k``: values, ``(B, L)`` for a single sub-column or ``(B, L, F)`` for a list;
+    - ``k_mask``: bool ``(B, L)``, True where a value is real rather than padding;
+    - ``k_lengths``: int64 ``(B,)``, the number of real values per row.
+
+    ``L`` is the longest row in the batch, after truncation. Rows whose nested value is
+    null or empty have length 0. Padding is zero, which carries no meaning; always use
+    the mask.
+
+    Parameters
+    ----------
+    column : str
+        The nested column.
+    fields : str or sequence of str
+        Sub-column name(s) to read.
+    dtype : torch.dtype, default torch.float32
+        Output dtype of the values.
+    max_length : int, optional
+        Longest sequence to keep; see ``truncate``.
+    truncate : {"raise", "head", "tail"}, default "raise"
+        What to do with rows longer than ``max_length``: raise, keep the first values, or
+        keep the last values.
+
+    Notes
+    -----
+    Missing values inside a list raise an error during collation.
+    """
+
+    column: str
+    fields: str | Sequence[str]
+    dtype: torch.dtype = torch.float32
+    max_length: int | None = None
+    truncate: TruncatePolicy = "raise"
+
+    def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None:
+        """Validate the nested column and sub-columns against ``meta``; see :meth:`Field.resolve`."""
+        subcolumns = _as_list(self.fields)
+        start = len(problems)
+        if not subcolumns:
+            problems.append(f"'{key}': Nested needs at least one sub-column.")
+        if self.truncate not in ("raise", "head", "tail"):
+            problems.append(f"'{key}': truncate must be 'raise', 'head' or 'tail', got {self.truncate!r}.")
+        limit = self.max_length
+        if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 1):
+            problems.append(f"'{key}': max_length must be a positive integer, got {limit!r}.")
+        if self.column not in meta.columns:
+            problems.append(f"'{key}': column '{self.column}' is not in the catalog.")
+        elif not isinstance(meta.dtypes[self.column], NestedDtype):
+            problems.append(f"'{key}': column '{self.column}' is not nested; use Flat('{self.column}').")
+        else:
+            available = meta.dtypes[self.column].column_dtypes
+            for sub in subcolumns:
+                if sub not in available:
+                    problems.append(
+                        f"'{key}': nested column '{self.column}' has no sub-column '{sub}' "
+                        f"(available: {sorted(available)})."
+                    )
+                elif not _is_numeric(available[sub]):
+                    problems.append(
+                        f"'{key}': sub-column '{self.column}.{sub}' has non-numeric type {available[sub]}."
+                    )
+        try:
+            np_dtype = numpy_dtype(self.dtype)
+        except TypeError as err:
+            problems.append(f"'{key}': {err}")
+        if len(problems) > start:
+            return None
+        return _ResolvedNested(key, subcolumns, isinstance(self.fields, str), self, np_dtype)
+
+
+class _ResolvedNested(ResolvedField):
+    def __init__(self, key: str, subcolumns: list[str], scalar: bool, spec: Nested, np_dtype: np.dtype):
+        self.key = key
+        self.columns = [spec.column]
+        self.subcolumns = subcolumns
+        self.scalar = scalar
+        self.spec = spec
+        self.np_dtype = np_dtype
+
+    def describe(self) -> dict[str, TensorInfo]:
+        values_shape = (None, None) if self.scalar else (None, None, len(self.subcolumns))
+        return {
+            self.key: TensorInfo(self.spec.dtype, values_shape),
+            f"{self.key}_mask": TensorInfo(torch.bool, (None, None)),
+            f"{self.key}_lengths": TensorInfo(torch.int64, (None,)),
+        }
+
+    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
+        what = f"'{self.key}' (nested column '{self.spec.column}', fields {self.subcolumns})"
+        series = frame[self.spec.column]
+        # nest.len() counts null rows as 0, and to_flat() returns the values of all rows
+        # concatenated in row order; both are consistent across iloc slices and concat.
+        lengths = np.asarray(series.nest.len(), dtype=np.int64)
+        values = frame_to_array(series.nest.to_flat(self.subcolumns), self.np_dtype, what)
+        padded, mask, kept = pad_ragged(
+            values, lengths, max_length=self.spec.max_length, truncate=self.spec.truncate, what=what
+        )
+        if self.scalar:
+            padded = padded[..., 0]
+        return {
+            self.key: torch.from_numpy(np.ascontiguousarray(padded)),
+            f"{self.key}_mask": torch.from_numpy(mask),
+            f"{self.key}_lengths": torch.from_numpy(kept),
+        }
+
+
+@dataclass
+class ColumnSpec:
+    """Mapping from output keys to catalog columns.
+
+    Parameters
+    ----------
+    fields : Mapping[str, Field]
+        Output key to field, e.g. ``{"x": Flat(["ra", "dec"]), "lc": Nested("lc", "flux")}``.
+    index_key : str or None, default "_index"
+        Every batch also includes the frame index (``_healpix_29`` for LSDB catalogs) as an
+        int64 ``(B,)`` tensor under this key, so batches can be traced back to rows. None
+        disables it.
+
+    Examples
+    --------
+    >>> import nested_pandas as npd
+    >>> nf = npd.NestedFrame({"a": [1.0, 2.0]}, index=[10, 11])
+    >>> nf = nf.join_nested(npd.NestedFrame({"f": [1.0, 2.0, 3.0]}, index=[10, 10, 11]), "lc")
+    >>> resolved = ColumnSpec({"a": Flat("a"), "lc": Nested("lc", "f")}).resolve(nf)
+    >>> batch = resolved.collate(nf)
+    >>> batch["lc"]
+    tensor([[1., 2.],
+            [3., 0.]])
+    >>> batch["lc_mask"]
+    tensor([[ True,  True],
+            [ True, False]])
+    >>> batch["_index"]
+    tensor([10, 11])
+    """
+
+    fields: Mapping[str, Field]
+    index_key: str | None = "_index"
+
+    def resolve(self, schema) -> ResolvedSpec:
+        """Validate the spec against a catalog's schema without computing any data.
+
+        Parameters
+        ----------
+        schema : lsdb.Catalog or pd.DataFrame
+            A catalog (its ``meta`` is used) or any dataframe with the same columns and
+            dtypes, such as ``catalog.meta`` or a computed partition.
+
+        Returns
+        -------
+        ResolvedSpec
+            The validated spec, which can collate frames.
+
+        Raises
+        ------
+        SpecError
+            Listing every problem found.
+        """
+        meta = schema.meta if hasattr(schema, "meta") and not isinstance(schema, pd.DataFrame) else schema
+        if not isinstance(meta, pd.DataFrame):
+            raise TypeError(f"Expected an lsdb catalog or a DataFrame, got {type(schema).__name__}.")
+        problems: list[str] = []
+        if not self.fields:
+            problems.append("The spec has no fields.")
+        resolved: dict[str, ResolvedField] = {}
+        for key, spec_field in self.fields.items():
+            if not isinstance(spec_field, Field):
+                problems.append(f"'{key}': expected a Field such as Flat or Nested, got {type(spec_field)}.")
+                continue
+            result = spec_field.resolve(key, meta, problems)
+            if result is not None:
+                resolved[key] = result
+
+        seen: dict[str, str] = {}
+        if self.index_key is not None:
+            seen[self.index_key] = "the index"
+            if not pd.api.types.is_integer_dtype(meta.index.dtype):
+                problems.append(
+                    f"index_key={self.index_key!r} needs an integer index, but the index has dtype "
+                    f"{meta.index.dtype}; pass index_key=None to disable it."
+                )
+        for key, result in resolved.items():
+            for output_key in result.describe():
+                if output_key in seen:
+                    problems.append(
+                        f"Output key '{output_key}' from '{key}' collides with {seen[output_key]}."
+                    )
+                seen[output_key] = f"'{key}'"
+
+        if problems:
+            hint = ""
+            missing = [p for p in problems if "is not in the catalog" in p]
+            all_columns = getattr(schema, "all_columns", None)
+            if missing and all_columns is not None:
+                hint = (
+                    "\nSome columns may exist on disk but were not loaded; the catalog was opened "
+                    f"with columns {list(meta.columns)}."
+                )
+            raise SpecError("Column spec does not match the catalog:\n  - " + "\n  - ".join(problems) + hint)
+        return ResolvedSpec(resolved, self.index_key)
+
+
+class ResolvedSpec:
+    """A :class:`ColumnSpec` validated against a schema; collates frames into tensors.
+
+    Create it with :meth:`ColumnSpec.resolve`. It is picklable, so it can be shipped to
+    Dask tasks.
+    """
+
+    def __init__(self, fields: dict[str, ResolvedField], index_key: str | None):
+        self.fields = fields
+        self.index_key = index_key
+
+    @property
+    def columns(self) -> list[str]:
+        """Top-level catalog columns needed for collation, in first-use order."""
+        columns: list[str] = []
+        for resolved in self.fields.values():
+            columns.extend(c for c in resolved.columns if c not in columns)
+        return columns
+
+    def describe(self) -> dict[str, TensorInfo]:
+        """Return the dtype and shape of every tensor a batch contains, keyed by output key."""
+        info: dict[str, TensorInfo] = {}
+        for resolved in self.fields.values():
+            info.update(resolved.describe())
+        if self.index_key is not None:
+            info[self.index_key] = TensorInfo(torch.int64, (None,))
+        return info
+
+    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
+        """Convert the rows of ``frame`` to a batch of tensors.
+
+        Parameters
+        ----------
+        frame : pd.DataFrame
+            Rows to convert, e.g. a slice of a catalog partition. Row order is preserved:
+            row ``i`` of every output tensor comes from row ``i`` of ``frame``.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            One entry per key in :meth:`describe`.
+        """
+        batch: dict[str, torch.Tensor] = {}
+        for resolved in self.fields.values():
+            batch.update(resolved.collate(frame))
+        if self.index_key is not None:
+            batch[self.index_key] = torch.from_numpy(np.array(frame.index.to_numpy(), dtype=np.int64))
+        return batch
diff --git a/tests/lsdb_ml/conftest.py b/tests/lsdb_ml/conftest.py
index e69de29..ec295e2 100644
--- a/tests/lsdb_ml/conftest.py
+++ b/tests/lsdb_ml/conftest.py
@@ -0,0 +1,57 @@
+import lsdb
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import pytest
+
+
+def make_nested_frame(n_rows: int = 500, seed: int = 0, max_length: int = 6) -> npd.NestedFrame:
+    """Synthetic object table with a ragged nested light curve column.
+
+    Light curve lengths range from 0 to ``max_length`` - 1, so some rows have empty
+    (null) nested values.
+    """
+    rng = np.random.default_rng(seed)
+    base = pd.DataFrame(
+        {
+            "id": np.arange(n_rows, dtype=np.int64),
+            "ra": rng.uniform(0, 360, n_rows),
+            "dec": rng.uniform(-60, 60, n_rows),
+            "mag": rng.normal(20, 1, n_rows).astype(np.float32),
+            "flag": rng.integers(0, 2, n_rows).astype(bool),
+            "name": [f"obj{i}" for i in range(n_rows)],
+        }
+    )
+    lengths = rng.integers(0, max_length, n_rows)
+    n_flat = int(lengths.sum())
+    flat = pd.DataFrame(
+        {
+            "t": rng.uniform(0, 100, n_flat),
+            "flux": rng.normal(0, 1, n_flat).astype(np.float32),
+            "band": rng.choice(["g", "r"], n_flat),
+        },
+        index=np.repeat(np.arange(n_rows), lengths),
+    )
+    return npd.NestedFrame(base).join_nested(flat, "lc")
+
+
+@pytest.fixture(scope="session")
+def nested_frame() -> npd.NestedFrame:
+    """The in-memory frame behind ``nested_catalog``, indexed by row number."""
+    return make_nested_frame()
+
+
+@pytest.fixture(scope="session")
+def nested_catalog(nested_frame) -> lsdb.Catalog:
+    """A lazy catalog with many small partitions and a ragged nested column."""
+    return lsdb.from_dataframe(nested_frame, partition_rows=40, catalog_name="toy")
+
+
+@pytest.fixture
+def unloadable_catalog(nested_catalog) -> lsdb.Catalog:
+    """Same schema as ``nested_catalog``, but computing any partition raises."""
+
+    def fail(_):
+        raise AssertionError("A partition was computed, but no data should be read here.")
+
+    return nested_catalog.map_partitions(fail, meta=nested_catalog.meta)
diff --git a/tests/lsdb_ml/test_collate.py b/tests/lsdb_ml/test_collate.py
new file mode 100644
index 0000000..b0123be
--- /dev/null
+++ b/tests/lsdb_ml/test_collate.py
@@ -0,0 +1,136 @@
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from lsdb_ml import ColumnSpec, Flat, Nested
+from lsdb_ml.collate import pad_ragged
+
+
+def reference_rows(frame, column, fields):
+    """Per-row python reference: a list of (length, n_fields) arrays."""
+    rows = []
+    for value in frame[column].array:
+        if value is None or not isinstance(value, pd.DataFrame):
+            rows.append(np.zeros((0, len(fields))))
+        else:
+            rows.append(value[fields].to_numpy(dtype=np.float64))
+    return rows
+
+
+def reference_pad(rows, width, keep="head"):
+    """Pad reference rows to ``width`` the slow way."""
+    padded = np.zeros((len(rows), width, rows[0].shape[1] if rows else 0))
+    mask = np.zeros((len(rows), width), dtype=bool)
+    for i, row in enumerate(rows):
+        row = row[:width] if keep == "head" else row[max(len(row) - width, 0) :]
+        padded[i, : len(row)] = row
+        mask[i, : len(row)] = True
+    return padded, mask
+
+
+@pytest.fixture
+def frame(nested_catalog):
+    """A frame mixing a sliced partition, a whole partition, and a concatenation."""
+    # A computed partition, plus a slice and a concatenation, all exercise the
+    # flat-values/offsets bookkeeping on non-trivial arrow layouts.
+    parts = [nested_catalog.partitions[i].compute(progress_bar=False) for i in range(3)]
+    return pd.concat([parts[0].iloc[3:], parts[1], parts[2].iloc[:5]])
+
+
+def test_flat_columns(frame):
+    """Flat columns stack into (B,) or (B, F) tensors with the requested dtype."""
+    resolved = ColumnSpec(
+        {"x": Flat(["ra", "dec", "mag"]), "mag": Flat("mag"), "flag": Flat("flag", dtype=torch.bool)}
+    ).resolve(frame)
+    batch = resolved.collate(frame)
+    assert batch["x"].shape == (len(frame), 3) and batch["x"].dtype == torch.float32
+    np.testing.assert_allclose(batch["x"].numpy(), frame[["ra", "dec", "mag"]].to_numpy(dtype=np.float32))
+    assert batch["mag"].shape == (len(frame),)
+    assert batch["flag"].dtype == torch.bool
+    np.testing.assert_array_equal(batch["flag"].numpy(), frame["flag"].to_numpy(dtype=bool))
+    np.testing.assert_array_equal(batch["_index"].numpy(), frame.index.to_numpy())
+    assert all(t.numpy().flags.writeable for t in batch.values())
+
+
+@pytest.mark.parametrize("fields", [["t", "flux"], ["flux"]])
+def test_nested_matches_reference(frame, fields):
+    """Vectorized nested padding matches a per-row reference, including empty rows."""
+    resolved = ColumnSpec({"lc": Nested("lc", fields, dtype=torch.float64)}).resolve(frame)
+    batch = resolved.collate(frame)
+    rows = reference_rows(frame, "lc", fields)
+    lengths = np.array([len(r) for r in rows])
+    assert (lengths == 0).any(), "fixture should include empty / null nested rows"
+    expected, expected_mask = reference_pad(rows, lengths.max())
+    np.testing.assert_allclose(batch["lc"].numpy(), expected)
+    np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask)
+    np.testing.assert_array_equal(batch["lc_lengths"].numpy(), lengths)
+
+
+def test_nested_single_field_is_2d(frame):
+    """A single nested sub-column produces (B, L) values."""
+    batch = ColumnSpec({"f": Nested("lc", "flux")}).resolve(frame).collate(frame)
+    assert batch["f"].ndim == 2
+    assert batch["f"].shape == batch["f_mask"].shape
+
+
+@pytest.mark.parametrize("keep", ["head", "tail"])
+def test_truncation(frame, keep):
+    """Rows over max_length keep their first or last values."""
+    resolved = ColumnSpec({"lc": Nested("lc", ["t"], max_length=2, truncate=keep, dtype=torch.float64)})
+    batch = resolved.resolve(frame).collate(frame)
+    rows = reference_rows(frame, "lc", ["t"])
+    expected, expected_mask = reference_pad(rows, 2, keep=keep)
+    np.testing.assert_allclose(batch["lc"].numpy(), expected)
+    np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask)
+    np.testing.assert_array_equal(batch["lc_lengths"].numpy(), np.minimum([len(r) for r in rows], 2))
+
+
+def test_truncation_raises_by_default(frame):
+    """Rows over max_length raise unless a truncate policy is given."""
+    with pytest.raises(ValueError, match="over max_length=2"):
+        ColumnSpec({"lc": Nested("lc", "t", max_length=2)}).resolve(frame).collate(frame)
+
+
+def test_empty_frame(frame):
+    """Zero-row frames collate to zero-row tensors."""
+    resolved = ColumnSpec({"x": Flat(["ra", "dec"]), "lc": Nested("lc", ["t", "flux"])}).resolve(frame)
+    batch = resolved.collate(frame.iloc[:0])
+    assert batch["x"].shape == (0, 2)
+    assert batch["lc"].shape == (0, 0, 2)
+    assert batch["lc_mask"].shape == (0, 0)
+    assert batch["_index"].shape == (0,)
+
+
+def test_all_rows_empty_nested(frame):
+    """A batch where every nested row is empty pads to width 0."""
+    lengths = np.asarray(frame["lc"].nest.len())
+    empty_rows = frame[lengths == 0]
+    batch = ColumnSpec({"lc": Nested("lc", "flux")}).resolve(frame).collate(empty_rows)
+    assert batch["lc"].shape == (len(empty_rows), 0)
+    assert not batch["lc_mask"].any()
+
+
+def test_null_values_raise():
+    """Missing flat values and missing values inside nested lists raise."""
+    nf = npd.NestedFrame({"a": [1.0, None, 3.0]}, index=[0, 1, 2])
+    nf = nf.join_nested(npd.NestedFrame({"f": [1.0, None, 2.0]}, index=[0, 0, 2]), "lc")
+    with pytest.raises(ValueError, match="1 row\\(s\\) with missing values"):
+        ColumnSpec({"a": Flat("a")}).resolve(nf).collate(nf)
+    with pytest.raises(ValueError, match="missing values"):
+        ColumnSpec({"lc": Nested("lc", "f")}).resolve(nf).collate(nf)
+
+
+def test_pad_ragged_length_mismatch():
+    """Inconsistent lengths and values are rejected."""
+    with pytest.raises(ValueError, match="lengths sum to 3"):
+        pad_ragged(np.zeros((2, 1)), np.array([1, 2]))
+
+
+def test_row_order_preserved(frame):
+    """Row i of every tensor comes from row i of the frame, even for a shuffled frame."""
+    shuffled = frame.sample(frac=1, random_state=1)
+    batch = ColumnSpec({"id": Flat("id", dtype=torch.int64)}).resolve(shuffled).collate(shuffled)
+    np.testing.assert_array_equal(batch["id"].numpy(), shuffled["id"].to_numpy())
+    np.testing.assert_array_equal(batch["_index"].numpy(), shuffled.index.to_numpy())
diff --git a/tests/lsdb_ml/test_spec.py b/tests/lsdb_ml/test_spec.py
new file mode 100644
index 0000000..1eb8faa
--- /dev/null
+++ b/tests/lsdb_ml/test_spec.py
@@ -0,0 +1,109 @@
+import pytest
+import torch
+
+from lsdb_ml import ColumnSpec, Flat, Nested, SpecError, TensorInfo
+
+
+def test_resolve_valid_spec_from_catalog_without_compute(unloadable_catalog):
+    """A valid spec resolves against a catalog without computing any partition."""
+    spec = ColumnSpec(
+        {
+            "x": Flat(["ra", "dec", "mag"]),
+            "flag": Flat("flag", dtype=torch.bool),
+            "lc": Nested("lc", ["t", "flux"], max_length=4, truncate="tail"),
+            "flux": Nested("lc", "flux", max_length=8, truncate="head"),
+        }
+    )
+    resolved = spec.resolve(unloadable_catalog)
+    assert resolved.columns == ["ra", "dec", "mag", "flag", "lc"]
+    assert resolved.describe() == {
+        "x": TensorInfo(torch.float32, (None, 3)),
+        "flag": TensorInfo(torch.bool, (None,)),
+        "lc": TensorInfo(torch.float32, (None, None, 2)),
+        "lc_mask": TensorInfo(torch.bool, (None, None)),
+        "lc_lengths": TensorInfo(torch.int64, (None,)),
+        "flux": TensorInfo(torch.float32, (None, None)),
+        "flux_mask": TensorInfo(torch.bool, (None, None)),
+        "flux_lengths": TensorInfo(torch.int64, (None,)),
+        "_index": TensorInfo(torch.int64, (None,)),
+    }
+
+
+def test_unloadable_catalog_fixture_raises(unloadable_catalog):
+    """Guard: the no-compute fixture really fails when a partition is computed."""
+    with pytest.raises(Exception, match="no data should be read"):
+        unloadable_catalog.partitions[0].compute(progress_bar=False)
+
+
+def test_resolve_accepts_meta_frame(nested_catalog):
+    """resolve() also accepts the meta frame directly."""
+    resolved = ColumnSpec({"x": Flat("ra")}).resolve(nested_catalog.meta)
+    assert resolved.columns == ["ra"]
+
+
+def test_all_problems_reported_together(unloadable_catalog):
+    """Every problem in a spec is reported in one error."""
+    spec = ColumnSpec(
+        {
+            "missing": Flat("nope"),
+            "nested_as_flat": Flat("lc"),
+            "string": Flat("name"),
+            "flat_as_nested": Nested("ra", "x"),
+            "missing_sub": Nested("lc", ["flux", "nope"]),
+            "string_sub": Nested("lc", "band"),
+            "bad_trunc": Nested("lc", "flux", truncate="middle"),
+            "bad_len": Nested("lc", "flux", max_length=0),
+        }
+    )
+    with pytest.raises(SpecError) as err:
+        spec.resolve(unloadable_catalog)
+    message = str(err.value)
+    for expected in [
+        "column 'nope' is not in the catalog",
+        "column 'lc' is nested",
+        "'name' has non-numeric dtype",
+        "column 'ra' is not nested",
+        "has no sub-column 'nope'",
+        "'lc.band' has non-numeric type",
+        "truncate must be",
+        "max_length must be a positive integer",
+    ]:
+        assert expected in message
+    assert "not loaded" in message  # hint for missing columns on a catalog
+
+
+def test_key_collisions(unloadable_catalog):
+    """Output keys may not collide with each other or the index key."""
+    with pytest.raises(SpecError, match="'lc_mask' from 'lc' collides"):
+        ColumnSpec({"lc_mask": Flat("ra"), "lc": Nested("lc", "flux")}).resolve(unloadable_catalog)
+    with pytest.raises(SpecError, match="'_index' from '_index' collides with the index"):
+        ColumnSpec({"_index": Flat("ra")}).resolve(unloadable_catalog)
+    ColumnSpec({"_index": Flat("ra")}, index_key=None).resolve(unloadable_catalog)
+
+
+def test_non_field_and_empty_spec(unloadable_catalog):
+    """Non-Field values and empty specs are rejected."""
+    with pytest.raises(SpecError, match="expected a Field"):
+        ColumnSpec({"x": "ra"}).resolve(unloadable_catalog)
+    with pytest.raises(SpecError, match="no fields"):
+        ColumnSpec({}).resolve(unloadable_catalog)
+
+
+def test_unsupported_dtype(unloadable_catalog):
+    """Torch dtypes without a numpy equivalent are rejected."""
+    with pytest.raises(SpecError, match="no numpy equivalent"):
+        ColumnSpec({"x": Flat("ra", dtype=torch.bfloat16)}).resolve(unloadable_catalog)
+
+
+def test_non_integer_index(nested_frame):
+    """The index key requires an integer index, and can be disabled."""
+    frame = nested_frame.reset_index(drop=True).set_index("name")
+    with pytest.raises(SpecError, match="needs an integer index"):
+        ColumnSpec({"x": Flat("ra")}).resolve(frame)
+    ColumnSpec({"x": Flat("ra")}, index_key=None).resolve(frame)
+
+
+def test_bad_schema_type():
+    """resolve() rejects objects that are neither catalogs nor frames."""
+    with pytest.raises(TypeError, match="Expected an lsdb catalog"):
+        ColumnSpec({"x": Flat("ra")}).resolve(42)

From 822a8c43835958e8c667c028479d2e08e0792b93 Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Thu, 15 Oct 2026 13:35:12 -0700
Subject: [PATCH 3/7] Add predict() with a per-worker model plugin

ModelDeployment registers a WorkerPlugin that builds the model once
per Dask worker from a ModelSpec (factory + state dict path, or an
in-memory module), in eval mode on a device resolved on the worker.
predict() maps each partition in order with task-local state and
attaches predictions by position, so output stays aligned with input
even for duplicate _healpix_29 values. OutputSpec builds the output
meta and checks every batch's output shape and row count.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/lsdb_ml/__init__.py                     |  22 +-
 src/lsdb_ml/inference/__init__.py           |   7 +
 src/lsdb_ml/inference/outputs.py            | 145 +++++++++++++
 src/lsdb_ml/inference/plugin.py             | 141 +++++++++++++
 src/lsdb_ml/inference/predict.py            | 212 ++++++++++++++++++++
 src/lsdb_ml/inputs.py                       |  41 ++++
 tests/lsdb_ml/inference/inference_models.py |  70 +++++++
 tests/lsdb_ml/inference/test_no_client.py   |  12 ++
 tests/lsdb_ml/inference/test_predict.py     | 171 ++++++++++++++++
 9 files changed, 820 insertions(+), 1 deletion(-)
 create mode 100644 src/lsdb_ml/inference/__init__.py
 create mode 100644 src/lsdb_ml/inference/outputs.py
 create mode 100644 src/lsdb_ml/inference/plugin.py
 create mode 100644 src/lsdb_ml/inference/predict.py
 create mode 100644 src/lsdb_ml/inputs.py
 create mode 100644 tests/lsdb_ml/inference/inference_models.py
 create mode 100644 tests/lsdb_ml/inference/test_no_client.py
 create mode 100644 tests/lsdb_ml/inference/test_predict.py

diff --git a/src/lsdb_ml/__init__.py b/src/lsdb_ml/__init__.py
index e42fde6..888e182 100644
--- a/src/lsdb_ml/__init__.py
+++ b/src/lsdb_ml/__init__.py
@@ -1,4 +1,24 @@
 from ._version import __version__
+from .device import move_batch, resolve_device
 from .example_module import greetings, meaning
+from .inference import ModelDeployment, ModelSpec, OutputSpec, predict
+from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo
 
-__all__ = ["greetings", "meaning", "__version__"]
+__all__ = [
+    "ColumnSpec",
+    "Field",
+    "Flat",
+    "Nested",
+    "ResolvedSpec",
+    "SpecError",
+    "TensorInfo",
+    "ModelDeployment",
+    "ModelSpec",
+    "OutputSpec",
+    "predict",
+    "move_batch",
+    "resolve_device",
+    "greetings",
+    "meaning",
+    "__version__",
+]
diff --git a/src/lsdb_ml/inference/__init__.py b/src/lsdb_ml/inference/__init__.py
new file mode 100644
index 0000000..9b696a4
--- /dev/null
+++ b/src/lsdb_ml/inference/__init__.py
@@ -0,0 +1,7 @@
+"""Running models over LSDB catalogs on Dask workers."""
+
+from .outputs import OutputSpec
+from .plugin import ModelPlugin, ModelSpec
+from .predict import ModelDeployment, predict
+
+__all__ = ["ModelDeployment", "ModelPlugin", "ModelSpec", "OutputSpec", "predict"]
diff --git a/src/lsdb_ml/inference/outputs.py b/src/lsdb_ml/inference/outputs.py
new file mode 100644
index 0000000..43c80a0
--- /dev/null
+++ b/src/lsdb_ml/inference/outputs.py
@@ -0,0 +1,145 @@
+"""Mapping model outputs to catalog columns."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import torch
+
+
+class OutputSpec:
+    """Names and dtypes of the columns a model's output becomes.
+
+    The model's output for a batch of ``B`` rows must be a tensor of shape ``(B,)`` or
+    ``(B, 1)`` for a single column, or ``(B, k)`` for ``k`` columns, in order.
+
+    Parameters
+    ----------
+    columns : Mapping[str, dtype] or sequence of str
+        Column name to pandas/numpy dtype (e.g. ``{"z": "float32"}``), or names that all
+        use ``dtype``.
+    dtype : default "float32"
+        The dtype for columns given as a sequence of names.
+
+    Examples
+    --------
+    >>> spec = OutputSpec(["p_star", "p_galaxy"])
+    >>> spec.columns
+    {'p_star': dtype('float32'), 'p_galaxy': dtype('float32')}
+    >>> spec.to_columns(torch.tensor([[0.1, 0.9], [0.8, 0.2]]), n_rows=2)["p_galaxy"]
+    array([0.9, 0.2], dtype=float32)
+    """
+
+    def __init__(self, columns: Mapping[str, object] | Sequence[str], dtype="float32"):
+        if isinstance(columns, str):
+            columns = [columns]
+        if not isinstance(columns, Mapping):
+            columns = dict.fromkeys(columns, dtype)
+        if not columns:
+            raise ValueError("OutputSpec needs at least one column.")
+        self.columns = {name: pd.api.types.pandas_dtype(d) for name, d in columns.items()}
+
+    def meta(self, input_meta: pd.DataFrame, append: bool) -> npd.NestedFrame:
+        """Return the empty output frame (Dask ``meta``) for an input catalog's meta.
+
+        Parameters
+        ----------
+        input_meta : pd.DataFrame
+            The input catalog's meta.
+        append : bool
+            Keep the input columns before the new ones.
+
+        Returns
+        -------
+        npd.NestedFrame
+            Empty frame with the output columns and dtypes.
+        """
+        empty = {name: pd.Series(dtype=d).array for name, d in self.columns.items()}
+        if append:
+            result = input_meta.copy()
+            for name, values in empty.items():
+                result[name] = values
+            return result
+        return npd.NestedFrame(empty, index=input_meta.index[:0])
+
+    def to_columns(self, output, n_rows: int) -> dict[str, np.ndarray]:
+        """Convert one batch of model output to numpy columns, validating the shape.
+
+        Parameters
+        ----------
+        output : torch.Tensor
+            Model output for a batch.
+        n_rows : int
+            Rows in the batch; every column must have exactly this many values.
+
+        Returns
+        -------
+        dict[str, np.ndarray]
+            One 1-d array per output column.
+
+        Raises
+        ------
+        ValueError
+            If the output does not match the spec.
+        """
+        names = list(self.columns)
+        if not isinstance(output, torch.Tensor):
+            raise TypeError(f"Expected the model output to be a torch.Tensor, got {type(output).__name__}.")
+        array = output.detach().to("cpu").numpy()
+        if array.ndim == 1:
+            array = array[:, None]
+        if array.ndim != 2 or array.shape[1] != len(names):
+            raise ValueError(
+                f"Model returned shape {tuple(array.shape)}, but {len(names)} output column(s) "
+                f"{names} were declared; expected ({n_rows}, {len(names)})."
+            )
+        arrays = {name: array[:, i] for i, name in enumerate(names)}
+        for name, array in arrays.items():
+            if len(array) != n_rows:
+                raise ValueError(
+                    f"Output '{name}' has {len(array)} values for a batch of {n_rows} rows; "
+                    "the model must return exactly one prediction per input row."
+                )
+        return arrays
+
+    def build_frame(
+        self, partition: pd.DataFrame, columns: dict[str, np.ndarray], append: bool
+    ) -> pd.DataFrame:
+        """Attach prediction columns to their partition, by position.
+
+        Parameters
+        ----------
+        partition : pd.DataFrame
+            The input partition the predictions were computed from, in the same row order.
+        columns : dict[str, np.ndarray]
+            One array per output column, ``len(partition)`` values each.
+        append : bool
+            Keep the input columns.
+
+        Returns
+        -------
+        pd.DataFrame
+            The output partition, with ``partition``'s index.
+        """
+        for name, values in columns.items():
+            if len(values) != len(partition):
+                raise ValueError(
+                    f"Output '{name}' has {len(values)} values for a partition of {len(partition)} rows."
+                )
+        # Assigning arrays (not Series) is positional: no index alignment, which would be
+        # wrong for a non-unique _healpix_29 index.
+        typed = {
+            name: pd.Series(columns[name], dtype=d, copy=False).array for name, d in self.columns.items()
+        }
+        if append:
+            result = partition.copy()
+            for name, values in typed.items():
+                result[name] = values
+            return result
+        result = npd.NestedFrame(index=partition.index)
+        for name, values in typed.items():
+            result[name] = values
+        return result
diff --git a/src/lsdb_ml/inference/plugin.py b/src/lsdb_ml/inference/plugin.py
new file mode 100644
index 0000000..fd0f131
--- /dev/null
+++ b/src/lsdb_ml/inference/plugin.py
@@ -0,0 +1,141 @@
+"""Installing a model once per Dask worker."""
+
+from __future__ import annotations
+
+import copy
+import hashlib
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import cloudpickle
+import torch
+from distributed import WorkerPlugin
+from torch import nn
+
+from lsdb_ml.device import DeviceLike, resolve_device
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    """How to build a model on a worker.
+
+    Only the spec travels to the workers: each worker calls ``factory()``, loads
+    ``weights`` itself, switches to ``eval()`` and moves the model to ``device``. Prefer a
+    factory plus a weights path for large models, so the weights never pass through the
+    Dask scheduler.
+
+    Parameters
+    ----------
+    factory : callable
+        Picklable ``factory() -> nn.Module``.
+    weights : str or Path, optional
+        A saved ``state_dict`` (``torch.save(model.state_dict(), path)``), readable from
+        every worker.
+    device : str or torch.device, default "cpu"
+        Resolved on each worker; ``"auto"`` picks CUDA, then MPS, then CPU.
+    torch_threads : int, optional
+        Calls ``torch.set_num_threads`` on the worker, to avoid oversubscribing cores with
+        Dask threads times torch intra-op threads. It affects the whole worker process.
+    """
+
+    factory: Callable[[], nn.Module]
+    weights: str | Path | None = None
+    device: DeviceLike = "cpu"
+    torch_threads: int | None = None
+
+    @classmethod
+    def from_module(cls, module: nn.Module, device: DeviceLike = "cpu") -> ModelSpec:
+        """Wrap an in-memory module; it is pickled into the plugin once per worker."""
+        return cls(factory=_CopyFactory(copy.deepcopy(module).to("cpu")), device=device)
+
+    def token(self) -> str:
+        """A short content hash, used to name the worker plugin."""
+        payload = cloudpickle.dumps((self.factory, str(self.weights), str(self.device), self.torch_threads))
+        return hashlib.sha256(payload).hexdigest()[:16]
+
+    def build(self) -> tuple[nn.Module, torch.device]:
+        """Build the model for inference: weights loaded, ``eval()`` mode, on its device."""
+        device = resolve_device(self.device)
+        if self.torch_threads is not None:
+            torch.set_num_threads(self.torch_threads)
+        model = self.factory()
+        if not isinstance(model, nn.Module):
+            raise TypeError(f"The model factory returned {type(model).__name__}, not a torch.nn.Module.")
+        if self.weights is not None:
+            model.load_state_dict(torch.load(self.weights, map_location=device, weights_only=True))
+        model.eval()
+        return model.to(device), device
+
+
+class _CopyFactory:
+    """Picklable factory returning a fresh copy of a module."""
+
+    def __init__(self, module: nn.Module):
+        self.module = module
+
+    def __call__(self) -> nn.Module:
+        return copy.deepcopy(self.module)
+
+
+class ModelPlugin(WorkerPlugin):
+    """Dask worker plugin that builds a model once when a worker starts (or joins).
+
+    Tasks look the model up with :func:`get_model`. Each worker has its own instance, so
+    there is no state shared between workers, and tasks on one worker share only the
+    read-only model.
+
+    Parameters
+    ----------
+    spec : ModelSpec
+        How to build the model.
+    name : str
+        Plugin name; tasks find the model by it.
+    """
+
+    def __init__(self, spec: ModelSpec, name: str):
+        self.spec = spec
+        self.name = name
+        self.model = None
+        self.device = None
+
+    def setup(self, worker=None):
+        """Build the model on this worker."""
+        self.model, self.device = self.spec.build()
+        logger.info("Loaded model plugin %s on %s (device %s).", self.name, worker, self.device)
+
+    def teardown(self, worker=None):
+        """Release the model."""
+        self.model = None
+        if self.device is not None and self.device.type == "cuda":
+            torch.cuda.empty_cache()
+
+
+def get_model(name: str) -> tuple[nn.Module, torch.device]:
+    """Return ``(model, device)`` for a deployed model, from inside a task.
+
+    Raises
+    ------
+    RuntimeError
+        If the model is not deployed where the task runs.
+    """
+    from distributed import get_worker
+
+    try:
+        worker = get_worker()
+    except ValueError as err:
+        raise RuntimeError(
+            f"Model '{name}' can only be used inside a Dask worker. Compute the prediction "
+            "catalog with the distributed Client it was deployed to."
+        ) from err
+    plugin = worker.plugins.get(name)
+    if plugin is None or plugin.model is None:
+        raise RuntimeError(
+            f"Model '{name}' is not deployed on worker {worker.address}. The prediction catalog "
+            "must be computed with the Dask client that was active (or passed) when predict() "
+            "was called, and before the deployment is closed."
+        )
+    return plugin.model, plugin.device
diff --git a/src/lsdb_ml/inference/predict.py b/src/lsdb_ml/inference/predict.py
new file mode 100644
index 0000000..1e917fd
--- /dev/null
+++ b/src/lsdb_ml/inference/predict.py
@@ -0,0 +1,212 @@
+"""Running a model over every row of an LSDB catalog."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Mapping, Sequence
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn
+
+from lsdb_ml.device import move_batch
+from lsdb_ml.inference.outputs import OutputSpec
+from lsdb_ml.inference.plugin import ModelPlugin, ModelSpec, get_model
+from lsdb_ml.inputs import InputsLike, model_args
+from lsdb_ml.spec import ColumnSpec, ResolvedSpec
+
+logger = logging.getLogger(__name__)
+
+# (client id, plugin name) pairs registered from this process, so redeploying the same
+# model to the same client does not rebuild it on every worker.
+_REGISTERED: set[tuple[str, str]] = set()
+
+
+def _active_client():
+    from distributed import get_client
+
+    try:
+        return get_client()
+    except ValueError as err:
+        raise RuntimeError(
+            "lsdb-ml inference needs a dask.distributed Client, e.g. Client(LocalCluster(...)) "
+            "or Client(processes=False); none is active."
+        ) from err
+
+
+def _iter_row_slices(frame: pd.DataFrame, batch_size: int) -> Iterator[pd.DataFrame]:
+    """Consecutive slices of at most ``batch_size`` rows, in order, covering ``frame`` once."""
+    for start in range(0, len(frame), batch_size):
+        yield frame.iloc[start : start + batch_size]
+
+
+class ModelDeployment:
+    """A model installed on every Dask worker, ready to predict on catalogs.
+
+    A :class:`~lsdb_ml.inference.plugin.ModelPlugin` builds the model once on each worker
+    of a ``dask.distributed`` client, including workers that join later.
+
+    Use it as a context manager to remove the model from the workers afterwards::
+
+        with ModelDeployment(ModelSpec(make_model, weights="model.pt")) as deployed:
+            result = deployed.predict(catalog, spec, inputs="x", outputs={"z": "float32"}).compute()
+
+    Parameters
+    ----------
+    model : ModelSpec or nn.Module
+        The model to deploy. A module is wrapped with :meth:`ModelSpec.from_module`.
+    client : distributed.Client, optional
+        Client to deploy to. Defaults to the active client.
+
+    Raises
+    ------
+    RuntimeError
+        If no client is given and none is active.
+    """
+
+    def __init__(self, model: ModelSpec | nn.Module, client=None):
+        self.spec = ModelSpec.from_module(model) if isinstance(model, nn.Module) else model
+        if not isinstance(self.spec, ModelSpec):
+            raise TypeError(f"Expected a ModelSpec or torch.nn.Module, got {type(model).__name__}.")
+        self.name = f"lsdb-ml-model-{self.spec.token()}"
+        self.client = client if client is not None else _active_client()
+        self._closed = False
+        if (self.client.id, self.name) not in _REGISTERED:
+            self.client.register_plugin(ModelPlugin(self.spec, self.name), name=self.name)
+            _REGISTERED.add((self.client.id, self.name))
+
+    def predict(
+        self,
+        catalog,
+        spec: ColumnSpec | ResolvedSpec,
+        outputs: OutputSpec | Mapping[str, object] | Sequence[str],
+        *,
+        inputs: InputsLike | None = None,
+        batch_size: int = 4096,
+        append: bool = True,
+    ):
+        """Return a lazy catalog with the model's predictions for every row.
+
+        See :func:`predict` for the parameters.
+        """
+        if self._closed:
+            raise RuntimeError("This deployment is closed.")
+        resolved = spec if isinstance(spec, ResolvedSpec) else spec.resolve(catalog)
+        if isinstance(batch_size, bool) or int(batch_size) < 1:
+            raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}.")
+        outputs = outputs if isinstance(outputs, OutputSpec) else OutputSpec(outputs)
+        if inputs is None:
+            if len(resolved.fields) != 1:
+                raise ValueError(
+                    f"The spec has several fields {list(resolved.fields)}; pass inputs= to say which "
+                    "batch keys the model takes."
+                )
+            inputs = next(iter(resolved.fields))
+        if append:
+            clashes = sorted(set(outputs.columns) & set(catalog.columns))
+            if clashes:
+                raise ValueError(f"Output columns {clashes} already exist in the catalog; rename them.")
+
+        meta = outputs.meta(catalog.meta, append)
+        return catalog.map_partitions(
+            _predict_partition,
+            model_name=self.name,
+            resolved=resolved,
+            inputs=inputs,
+            outputs=outputs,
+            batch_size=int(batch_size),
+            append=append,
+            meta=meta,
+        )
+
+    def close(self) -> None:
+        """Remove the model from the workers. Safe to call twice."""
+        if self._closed:
+            return
+        self._closed = True
+        self.client.unregister_worker_plugin(self.name)
+        _REGISTERED.discard((self.client.id, self.name))
+
+    def __enter__(self) -> ModelDeployment:
+        return self
+
+    def __exit__(self, *exc_info) -> None:
+        self.close()
+
+
+def predict(
+    catalog,
+    model: ModelSpec | nn.Module,
+    spec: ColumnSpec | ResolvedSpec,
+    outputs: OutputSpec | Mapping[str, object] | Sequence[str],
+    *,
+    inputs: InputsLike | None = None,
+    batch_size: int = 4096,
+    append: bool = True,
+    client=None,
+):
+    """Run a model over every row of a catalog, returning a lazy catalog of predictions.
+
+    The model is installed once per Dask worker (see :class:`ModelDeployment`) and stays
+    installed so the returned catalog can be computed later; use
+    :class:`ModelDeployment` directly to remove it when done.
+
+    Each partition is processed in order, in batches of up to ``batch_size`` rows, under
+    ``torch.inference_mode()``. Predictions are attached to the rows they were computed
+    from by position within the partition, never by joining on the index, so the output
+    is aligned with the input even when ``_healpix_29`` values repeat.
+
+    Parameters
+    ----------
+    catalog : lsdb.Catalog
+        The input catalog.
+    model : ModelSpec or nn.Module
+        The model, or how to build it on each worker.
+    spec : ColumnSpec or ResolvedSpec
+        Which columns become which tensors. Validated against ``catalog.meta`` now,
+        before any data is read.
+    outputs : OutputSpec, mapping or sequence
+        Output column names and dtypes, e.g. ``{"z_phot": "float32"}``.
+    inputs : str, sequence of str or callable, optional
+        Batch key(s) passed positionally to the model, or ``callable(batch)``. Defaults to
+        the spec's only field.
+    batch_size : int, default 4096
+        Maximum rows per forward pass.
+    append : bool, default True
+        Keep the input columns. With False, the result holds only the predictions (on the
+        same index).
+    client : distributed.Client, optional
+        Client to deploy to. Defaults to the active client; one is required.
+
+    Returns
+    -------
+    lsdb.Catalog
+        Lazy catalog whose ``meta`` includes the output columns.
+    """
+    deployment = ModelDeployment(model, client=client)
+    return deployment.predict(catalog, spec, outputs, inputs=inputs, batch_size=batch_size, append=append)
+
+
+def _predict_partition(
+    partition: pd.DataFrame,
+    *,
+    model_name: str,
+    resolved: ResolvedSpec,
+    inputs: InputsLike,
+    outputs: OutputSpec,
+    batch_size: int,
+    append: bool,
+) -> pd.DataFrame:
+    """Predict on one partition. All state is local to this call except the read-only model."""
+    if len(partition) == 0:
+        empty = {name: np.empty(0) for name in outputs.columns}
+        return outputs.build_frame(partition, empty, append)
+    model, device = get_model(model_name)
+    pieces: list[dict[str, np.ndarray]] = []
+    with torch.inference_mode():
+        for rows in _iter_row_slices(partition, batch_size):
+            batch = move_batch(resolved.collate(rows), device)
+            pieces.append(outputs.to_columns(model(*model_args(batch, inputs)), len(rows)))
+    columns = {name: np.concatenate([piece[name] for piece in pieces]) for name in outputs.columns}
+    return outputs.build_frame(partition, columns, append)
diff --git a/src/lsdb_ml/inputs.py b/src/lsdb_ml/inputs.py
new file mode 100644
index 0000000..6e18329
--- /dev/null
+++ b/src/lsdb_ml/inputs.py
@@ -0,0 +1,41 @@
+"""Selecting model inputs from a collated batch."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+import torch
+
+#: How to build the model's positional arguments from a batch: one key, a list of keys,
+#: or a callable ``batch -> argument`` (a tuple return value is unpacked into arguments).
+InputsLike = str | Sequence[str] | Callable[[Mapping[str, torch.Tensor]], Any]
+
+
+def model_args(batch: Mapping[str, torch.Tensor], inputs: InputsLike) -> tuple:
+    """Return the positional arguments to call a model with.
+
+    Parameters
+    ----------
+    batch : Mapping[str, torch.Tensor]
+        A collated batch.
+    inputs : str, sequence of str, or callable
+        A key, keys (passed positionally in order), or ``callable(batch)``.
+
+    Returns
+    -------
+    tuple
+        Positional arguments.
+
+    Examples
+    --------
+    >>> batch = {"x": 1, "y": 2}
+    >>> model_args(batch, "x"), model_args(batch, ["y", "x"]), model_args(batch, lambda b: (b["x"], 5))
+    ((1,), (2, 1), (1, 5))
+    """
+    if isinstance(inputs, str):
+        return (batch[inputs],)
+    if callable(inputs):
+        result = inputs(batch)
+        return result if isinstance(result, tuple) else (result,)
+    return tuple(batch[key] for key in inputs)
diff --git a/tests/lsdb_ml/inference/inference_models.py b/tests/lsdb_ml/inference/inference_models.py
new file mode 100644
index 0000000..3364717
--- /dev/null
+++ b/tests/lsdb_ml/inference/inference_models.py
@@ -0,0 +1,70 @@
+"""Tiny models and checks shared by the inference tests."""
+
+import numpy as np
+import torch
+from torch import nn
+
+from lsdb_ml import ColumnSpec, Flat
+
+# Counts factory calls; with processes=False every worker shares this module.
+BUILDS = {"count": 0}
+
+
+class AffineMag(nn.Module):
+    """pred = 2 * mag + 1, exactly, so every row can be checked against its own input."""
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(1, 1, dtype=torch.float64)
+        with torch.no_grad():
+            self.linear.weight.fill_(2.0)
+            self.linear.bias.fill_(1.0)
+
+    def forward(self, mag):
+        """Map (B,) magnitudes to (B,) predictions."""
+        return self.linear(mag.unsqueeze(-1)).squeeze(-1)
+
+
+def counting_factory():
+    """Build AffineMag and count how often a model is built."""
+    BUILDS["count"] += 1
+    return AffineMag()
+
+
+class MaskedMean(nn.Module):
+    """Mean of the real (unpadded) light-curve values; 0 for empty light curves."""
+
+    def forward(self, flux, mask):
+        """Average padded values using the mask."""
+        return (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+
+
+class Stacked(nn.Module):
+    """Two outputs per row, as a (B, 2) tensor."""
+
+    def forward(self, mag):
+        """Return [mag, -mag]."""
+        return torch.stack([mag, -mag], dim=1)
+
+
+class DropsRow(nn.Module):
+    """A broken model that returns one prediction too few."""
+
+    def forward(self, mag):
+        """Drop the last row."""
+        return mag[:-1]
+
+
+MAG_SPEC = ColumnSpec({"mag": Flat("mag", dtype=torch.float64)})
+
+
+def check_alignment(result, source_frame):
+    """Every prediction matches its own row's input, and rows match the source exactly once."""
+    np.testing.assert_allclose(
+        result["pred"].to_numpy(), 2 * result["mag"].to_numpy(np.float64) + 1, rtol=1e-6
+    )
+    # Also join on the object id against an eager computation on the raw frame.
+    expected = 2 * source_frame.set_index("id")["mag"].astype(np.float64) + 1
+    assert result["id"].is_unique
+    joined = result.set_index("id")["pred"]
+    np.testing.assert_allclose(joined.to_numpy(), expected.loc[joined.index].to_numpy(), rtol=1e-6)
diff --git a/tests/lsdb_ml/inference/test_no_client.py b/tests/lsdb_ml/inference/test_no_client.py
new file mode 100644
index 0000000..bb8960f
--- /dev/null
+++ b/tests/lsdb_ml/inference/test_no_client.py
@@ -0,0 +1,12 @@
+"""Behavior without a distributed client; kept apart from the module-scoped client fixture."""
+
+import pytest
+from inference_models import AffineMag
+
+from lsdb_ml import ModelDeployment
+
+
+def test_requires_distributed_client():
+    """Without an active client, deploying fails with a clear message."""
+    with pytest.raises(RuntimeError, match="needs a dask.distributed Client"):
+        ModelDeployment(AffineMag())
diff --git a/tests/lsdb_ml/inference/test_predict.py b/tests/lsdb_ml/inference/test_predict.py
new file mode 100644
index 0000000..554616a
--- /dev/null
+++ b/tests/lsdb_ml/inference/test_predict.py
@@ -0,0 +1,171 @@
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+from distributed import Client, LocalCluster
+from inference_models import (
+    BUILDS,
+    MAG_SPEC,
+    AffineMag,
+    DropsRow,
+    MaskedMean,
+    Stacked,
+    check_alignment,
+    counting_factory,
+)
+
+from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested, OutputSpec, predict
+
+
+@pytest.fixture(scope="module")
+def client():
+    """Two in-process workers with two threads each, so tasks run concurrently."""
+    with (
+        LocalCluster(n_workers=2, threads_per_worker=2, processes=False, dashboard_address=None) as cluster,
+        Client(cluster) as dask_client,
+    ):
+        yield dask_client
+
+
+@pytest.fixture
+def scrambled_catalog(nested_catalog):
+    """A non-trivial layout: a filter that empties some partitions, then a reordered subset."""
+    filtered = nested_catalog.query("mag > 21")
+    n = filtered.npartitions
+    return filtered.partitions[[n - 1, 3, 0, 7, 5, 2, 9, 11, 1, 4, 6, 8, 10] + list(range(12, n - 1))]
+
+
+def test_index_aligned_after_nontrivial_layout(client, scrambled_catalog, nested_frame):
+    """Predictions line up with their input rows across empty, reordered partitions."""
+    computed_input = scrambled_catalog.compute(progress_bar=False)
+    assert 0 in [len(p.compute(progress_bar=False)) for p in scrambled_catalog.partitions]
+
+    with ModelDeployment(ModelSpec(AffineMag)) as deployed:
+        lazy = deployed.predict(scrambled_catalog, MAG_SPEC, {"pred": "float64"}, batch_size=3)
+        result = lazy.compute(progress_bar=False)
+    assert len(result) == len(computed_input)
+    assert result.index.equals(computed_input.index)
+    check_alignment(result, nested_frame)
+
+
+def test_alignment_with_duplicate_index(client, nested_frame):
+    """Rows sharing a _healpix_29 value still get their own predictions."""
+    import lsdb
+
+    frame = nested_frame.iloc[:60].copy()
+    frame["ra"] = 10.0  # every object at the same position: one index value repeated
+    frame["dec"] = 10.0
+    catalog = lsdb.from_dataframe(frame, partition_rows=100)
+    result = predict(catalog, AffineMag(), MAG_SPEC, {"pred": "float64"}, batch_size=7).compute(
+        progress_bar=False
+    )
+    assert not result.index.is_unique
+    check_alignment(result, nested_frame)
+
+
+def test_model_loaded_once_per_worker(client, nested_catalog):
+    """The worker plugin builds the model once per worker, not once per partition."""
+    BUILDS["count"] = 0
+    spec = ModelSpec(counting_factory)
+    with ModelDeployment(spec) as deployed:
+        deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+        deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+        # Deploying the same spec again reuses the installed plugin.
+        predict(nested_catalog, spec, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+    assert nested_catalog.npartitions > len(client.scheduler_info()["workers"])
+    assert BUILDS["count"] == len(client.scheduler_info()["workers"]) == 2
+
+
+def test_meta_matches_computed(client, nested_catalog):
+    """The lazy result's meta has the same columns and dtypes as the computed data."""
+    for append in (True, False):
+        lazy = predict(nested_catalog, AffineMag(), MAG_SPEC, {"pred": "float32"}, append=append)
+        computed = lazy.compute(progress_bar=False)
+        assert list(lazy.columns) == list(computed.columns)
+        assert lazy.dtypes.to_dict() == computed.dtypes.to_dict()
+        assert computed["pred"].dtype == np.float32
+    assert list(lazy.columns) == ["pred"]
+    assert len(computed) == 500
+
+
+def test_nested_inputs(client, nested_catalog):
+    """Nested columns are padded and masked, so a masked mean matches a per-row mean."""
+    spec = ColumnSpec({"flux": Nested("lc", "flux", dtype=torch.float64)})
+    result = predict(
+        nested_catalog,
+        MaskedMean(),
+        spec,
+        {"mean_flux": "float64"},
+        inputs=["flux", "flux_mask"],
+        batch_size=9,
+    ).compute(progress_bar=False)
+    expected = [
+        float(np.mean(v["flux"])) if isinstance(v, pd.DataFrame) and len(v) else 0.0
+        for v in result["lc"].array
+    ]
+    np.testing.assert_allclose(result["mean_flux"].to_numpy(), expected, rtol=1e-6)
+
+
+def test_multi_column_output(client, nested_catalog):
+    """A (B, k) output fills k columns in order."""
+    result = predict(nested_catalog, Stacked(), MAG_SPEC, ["plus", "minus"], append=False).compute(
+        progress_bar=False
+    )
+    np.testing.assert_allclose(result["plus"], -result["minus"])
+
+
+def test_wrong_row_count_fails(client, nested_catalog):
+    """A model that drops rows fails loudly instead of misaligning predictions."""
+    lazy = predict(nested_catalog, DropsRow(), MAG_SPEC, {"pred": "float32"})
+    with pytest.raises(Exception, match="exactly one prediction per input row"):
+        lazy.compute(progress_bar=False)
+
+
+def test_validation(client, unloadable_catalog):
+    """Bad specs, outputs and arguments fail before any data is read."""
+    with pytest.raises(ValueError, match="already exist"):
+        predict(unloadable_catalog, AffineMag(), MAG_SPEC, {"mag": "float32"})
+    with pytest.raises(ValueError, match="pass inputs="):
+        predict(unloadable_catalog, AffineMag(), ColumnSpec({"a": Flat("ra"), "b": Flat("dec")}), ["p"])
+    with pytest.raises(ValueError, match="not in the catalog"):
+        predict(unloadable_catalog, AffineMag(), ColumnSpec({"a": Flat("nope")}), ["p"])
+    with pytest.raises(ValueError, match="batch_size"):
+        predict(unloadable_catalog, AffineMag(), MAG_SPEC, ["p"], batch_size=0)
+    with pytest.raises(TypeError, match="ModelSpec or torch.nn.Module"):
+        ModelDeployment("model")
+
+
+def test_closed_deployment(client, nested_catalog):
+    """After close, the model is gone from the workers and computing fails clearly."""
+    deployed = ModelDeployment(ModelSpec(AffineMag, device="cpu", torch_threads=None))
+    lazy = deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float32"})
+    deployed.close()
+    deployed.close()
+    with pytest.raises(Exception, match="is not deployed"):
+        lazy.compute(progress_bar=False)
+    with pytest.raises(RuntimeError, match="closed"):
+        deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float32"})
+
+
+def test_weights_loaded_on_workers(client, nested_catalog, tmp_path):
+    """ModelSpec loads a saved state dict on each worker."""
+    trained = AffineMag()
+    with torch.no_grad():
+        trained.linear.bias.fill_(5.0)
+    torch.save(trained.state_dict(), tmp_path / "state.pt")
+    spec = ModelSpec(AffineMag, weights=tmp_path / "state.pt")
+    result = predict(nested_catalog, spec, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+    np.testing.assert_allclose(result["pred"], 2 * result["mag"].astype(np.float64) + 5, rtol=1e-6)
+
+
+def test_output_spec_errors():
+    """Outputs that do not match the declared columns are rejected."""
+    spec = OutputSpec({"a": "float32", "b": "float32"})
+    with pytest.raises(ValueError, match="expected \\(3, 2\\)"):
+        spec.to_columns(torch.zeros(3), 3)
+    with pytest.raises(TypeError, match="torch.Tensor"):
+        spec.to_columns([1, 2, 3], 3)
+    with pytest.raises(ValueError, match="at least one column"):
+        OutputSpec([])
+    assert set(spec.to_columns(torch.zeros(3, 2), 3)) == {"a", "b"}
+    assert OutputSpec(["z"]).to_columns(torch.zeros(3, 1), 3)["z"].shape == (3,)

From 7ae5f6249ea15f3f9c9b2b30e609e00a002d43ad Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Thu, 15 Oct 2026 13:35:12 -0700
Subject: [PATCH 4/7] Add inference example script

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/examples/predict_on_catalog.py | 80 +++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)
 create mode 100644 docs/examples/predict_on_catalog.py

diff --git a/docs/examples/predict_on_catalog.py b/docs/examples/predict_on_catalog.py
new file mode 100644
index 0000000..b68b051
--- /dev/null
+++ b/docs/examples/predict_on_catalog.py
@@ -0,0 +1,80 @@
+"""Run a PyTorch model over every row of an LSDB catalog on Dask workers.
+
+Run with ``python docs/examples/predict_on_catalog.py``. It uses a synthetic catalog;
+for a real one, replace ``make_catalog()`` with ``lsdb.open_catalog(...)``.
+
+The model is built once per Dask worker (by a worker plugin) and the result is a lazy
+catalog with prediction columns, aligned row by row with the input.
+"""
+
+from pathlib import Path
+
+import lsdb
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import torch
+from distributed import Client, LocalCluster
+from torch import nn
+
+from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested
+
+
+def make_catalog(n_objects=5000, seed=0):
+    """Objects with a magnitude and a ragged light curve."""
+    rng = np.random.default_rng(seed)
+    lengths = rng.integers(0, 40, n_objects)
+    objects = pd.DataFrame(
+        {
+            "object_id": np.arange(n_objects),
+            "ra": rng.uniform(0, 360, n_objects),
+            "dec": rng.uniform(-60, 60, n_objects),
+            "mag": rng.normal(20, 1, n_objects).astype(np.float32),
+        }
+    )
+    light_curves = pd.DataFrame(
+        {"mjd": rng.uniform(0, 100, lengths.sum()), "flux": rng.normal(1, 0.1, lengths.sum())},
+        index=np.repeat(np.arange(n_objects), lengths),
+    )
+    return lsdb.from_dataframe(npd.NestedFrame(objects).join_nested(light_curves, "lc"), partition_rows=500)
+
+
+class Scorer(nn.Module):
+    """Toy model combining a flat column with a masked summary of a nested column."""
+
+    def __init__(self):
+        super().__init__()
+        self.head = nn.Linear(2, 2)
+
+    def forward(self, mag, flux, mask):
+        """Return two scores per object from the magnitude and the mean light-curve flux."""
+        mean_flux = (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        return self.head(torch.stack([mag, mean_flux], dim=1))
+
+
+def main(weights_path):
+    """Save some weights, deploy the model and predict on the catalog."""
+    torch.manual_seed(0)
+    torch.save(Scorer().state_dict(), weights_path)  # stands in for a trained model
+
+    catalog = make_catalog().query("mag < 21")  # any lazy lsdb pipeline works here
+    spec = ColumnSpec({"mag": Flat("mag"), "flux": Nested("lc", "flux", max_length=32, truncate="tail")})
+
+    # Workers build Scorer() and load the weights themselves; only the spec is shipped.
+    with ModelDeployment(ModelSpec(Scorer, weights=weights_path, device="cpu")) as deployed:
+        scored = deployed.predict(
+            catalog,
+            spec,
+            outputs={"score_a": "float32", "score_b": "float32"},
+            inputs=["mag", "flux", "flux_mask"],
+            batch_size=1024,
+        )
+        print(scored.dtypes[["object_id", "mag", "score_a", "score_b"]])  # meta, nothing computed yet
+        result = scored.compute()
+    print(result[["object_id", "mag", "score_a", "score_b"]].head())
+    print(f"{len(result)} rows scored")
+
+
+if __name__ == "__main__":
+    with LocalCluster(n_workers=2, threads_per_worker=2, dashboard_address=None) as cluster, Client(cluster):
+        main(Path("scorer_weights.pt").resolve())

From e3915726108343eeecc33d0752c80382decc531a Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Thu, 15 Oct 2026 13:37:14 -0700
Subject: [PATCH 5/7] Add MVP design doc

Trimmed design covering only what this branch ships (column spec and
inference), with deferred items pointing to the full design.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/design/mvp-design.md | 138 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)
 create mode 100644 docs/design/mvp-design.md

diff --git a/docs/design/mvp-design.md b/docs/design/mvp-design.md
new file mode 100644
index 0000000..c8d23a9
--- /dev/null
+++ b/docs/design/mvp-design.md
@@ -0,0 +1,138 @@
+# lsdb-ml MVP: column spec and inference
+
+Status: implemented on this branch. This document covers only what the MVP ships. The fuller
+design, which also covers training, is on the `design/initial-design` branch; the parts of it
+that this MVP leaves out are listed in §7.
+
+## 1. Scope
+
+This MVP runs a PyTorch model over every row of an LSDB catalog on Dask workers. It returns
+a lazy catalog of predictions that has correct `meta` and is aligned row by row with the
+input. Two pieces make that up:
+
+- **Column spec:** maps flat and nested catalog columns to tensors, validated against
+  catalog `meta` before any data is read.
+- **Inference:** a Dask `WorkerPlugin` installs the model once per worker, and
+  `map_partitions` runs it partition by partition.
+
+Training, the Hyrax backend, and running without a `distributed` client are out of scope.
+`hyrax` remains only an optional extra.
+
+## 2. API
+
+```python
+import torch
+from distributed import Client
+from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested
+
+client = Client(...)  # a dask.distributed client is required
+
+spec = ColumnSpec({
+    "mag": Flat("mag"),                                              # -> (B,)
+    "flux": Nested("lc", "flux", max_length=32, truncate="tail"),    # -> (B, L), flux_mask, flux_lengths
+})
+
+with ModelDeployment(ModelSpec(MyModel, weights="model.pt")) as deployed:
+    scored = deployed.predict(
+        catalog, spec,
+        outputs={"score": "float32"},
+        inputs=["mag", "flux", "flux_mask"],   # batch keys passed positionally to the model
+        batch_size=4096,
+    )
+    result = scored.compute()
+```
+
+`predict(catalog, model, spec, outputs, ...)` is the one-call form of the same thing. It
+leaves the model installed, because the returned catalog is lazy. `ModelDeployment.close()`
+removes the model from the workers.
+
+## 3. Column spec
+
+- **Fields.** `Flat(columns, dtype=float32)` stacks numeric columns into a `(B,)` or
+  `(B, F)` tensor. `Nested(column, fields, dtype=float32, max_length=None, truncate="raise")`
+  reads numeric sub-columns of a nested column.
+- **Validation before compute.** `ColumnSpec.resolve(catalog)` uses only `catalog.meta`.
+  Meta is an empty `NestedFrame` with full dtypes, and `NestedDtype.column_dtypes` gives the
+  sub-column types. That replaces the hand-built column preview in hyrax#1062. Resolve checks:
+  - columns and sub-columns exist;
+  - flat vs nested is used correctly;
+  - dtypes are numeric;
+  - output keys don't collide.
+
+  Every problem is reported in one `SpecError`.
+- **Nested representation: zero padding with a mask.** Each batch is padded to its longest
+  row and returns `<key>` `(B, L[, F])`, `<key>_mask` `(B, L)` and `<key>_lengths` `(B,)`.
+  This works directly with attention masks, masked pooling and `pack_padded_sequence`.
+  `torch.nested` is not mature enough to be the default. Other representations can be added
+  later as new `Field` subclasses.
+- **Collation is vectorized.** It is built from `series.nest.len()` (null rows count as 0)
+  and `series.nest.to_flat()`, with no per-row Python.
+- **Missing values raise.** Drop or fill them in the catalog first.
+- **Every batch includes `_index`.** It holds the frame index (`_healpix_29`) as int64.
+
+## 4. Inference lifecycle
+
+- **Model spec.** `ModelSpec(factory, weights=None, device="cpu", torch_threads=None)`. Only
+  the factory and the weights path go to the scheduler. Each worker builds the model, loads
+  the `state_dict`, calls `eval()` and moves it to the device. The device is resolved *on the
+  worker*: `"auto"` picks CUDA, then MPS, then CPU. A bare `nn.Module` is also accepted; it is
+  pickled into the plugin once per worker.
+- **Plugin.** `ModelPlugin(WorkerPlugin)` is named `lsdb-ml-model-<hash of spec>`. Concurrent
+  deployments therefore don't collide, and redeploying the same spec to the same client
+  doesn't rebuild the model. Workers that join later get the plugin automatically.
+- **Tasks.** Each task looks the model up with `get_worker().plugins[name]`. If the model is
+  missing, the task raises a clear error, e.g. when the catalog is computed on another client
+  or after `close()`.
+- **Concurrency.** Tasks on one worker share only the read-only model, under
+  `torch.inference_mode()`. GPU workers should use `threads_per_worker=1`.
+
+## 5. Alignment (the hyrax#1062 "shuffled results" issue)
+
+The prototype kept one mutable dataset per worker, overwrote it in each task, and wrote
+predictions back by position into a buffer. Concurrent tasks on a worker could race on that
+shared state. Here, each task:
+
+1. holds only task-local state;
+2. walks its partition in order, in `iloc` slices of `batch_size` rows;
+3. checks every batch's output for exactly one row per input row;
+4. assigns the concatenated predictions **by position** to the same partition frame they came
+   from.
+
+Predictions are never joined on `_healpix_29`, which is not guaranteed unique. The output
+`meta` comes from `catalog.meta` plus the `OutputSpec` dtypes. Columns are built with
+`pd.Series(values, dtype=d).array`, so computed dtypes match the meta exactly.
+
+## 6. Tests
+
+All tests run on CPU, take about 16 seconds, and use a synthetic nested catalog and an
+in-process 2-worker × 2-thread cluster.
+
+- **Spec validation:** every error case, including a fixture catalog that raises if any
+  partition is computed.
+- **Collation:** checked against a per-row reference, covering truncation, empty and null
+  nested rows, missing values, and row order.
+- **Alignment:** a filter that empties partitions, then a reordered partition subset; and a
+  catalog where every row shares one `_healpix_29` value. Both are checked per row and joined
+  on object id. Injecting a reversed batch order makes these tests fail.
+- **Load once per worker:** the model is built exactly once per worker across many
+  partitions and repeated predictions.
+- **Meta:** dtypes of the lazy result match the computed result.
+- **Failure modes:** wrong output row count, closed deployment, no client, bad arguments.
+
+## 7. Deferred (in the full design)
+
+| Item | Why deferred |
+|---|---|
+| Training: `CatalogIterableDataset` (partition sharding across DataLoader workers and DDP) and `TrainSession` | Larger surface; depends on upstream `CatalogStream` changes. Planned as two follow-up PRs. |
+| `transform=` hooks, `pad_to`/`pad_value`, missing-value filling | Convenience; easy to add back when needed. |
+| Dict-valued and list-column model outputs | Not needed for the first use cases. |
+| Running without a `distributed` client | A second code path; requiring a `Client` is a reasonable MVP constraint. |
+| Hyrax backend | The seams already exist: the model factory in `ModelSpec` and `inputs=` (a key list or callable). |
+| Automatic output meta from a dummy forward pass | Fragile for models that can't take an empty or synthetic batch. |
+
+## 8. CI note
+
+The workflows only trigger on `main`, so PRs into `dev` get no CI. Add `dev` to the
+`push`/`pull_request` branch lists in `testing-and-coverage.yml`, `pre-commit-ci.yml` and
+`build-documentation.yml`. With torch as a dependency, CI should also install the CPU build
+of torch, because Linux defaults to the multi-GB CUDA wheels.

From ebb672abf4088eeed868a7a8e6d22ca86835b293 Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Fri, 16 Oct 2026 14:37:55 -0700
Subject: [PATCH 6/7] Replace ColumnSpec with column-name inputs inferred from
 meta

Inputs are now plain catalog column names: flat columns by name and
nested sub-columns by LSDB's dotted names ('lc.flux'). Flat vs nested
and dtypes are read from catalog meta, checked before any data is
read, and dry-run on the empty meta frame. Batches are keyed by those
names; each nested column contributes a shared '<column>._mask' and
'<column>._lengths'. The model receives the batch dict, or a
forward(model, batch) adapter for positional models. Data preparation
is left to ordinary LSDB operations.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/design/mvp-design.md                   | 106 +++--
 docs/examples/predict_on_catalog.py         |  15 +-
 src/lsdb_ml/__init__.py                     |  14 +-
 src/lsdb_ml/batch.py                        | 264 ++++++++++++
 src/lsdb_ml/collate.py                      |  34 +-
 src/lsdb_ml/inference/predict.py            |  87 ++--
 src/lsdb_ml/inputs.py                       |  41 --
 src/lsdb_ml/spec.py                         | 423 --------------------
 tests/lsdb_ml/inference/inference_models.py |  23 +-
 tests/lsdb_ml/inference/test_predict.py     |  58 +--
 tests/lsdb_ml/test_batch.py                 | 181 +++++++++
 tests/lsdb_ml/test_collate.py               | 136 -------
 tests/lsdb_ml/test_spec.py                  | 109 -----
 13 files changed, 628 insertions(+), 863 deletions(-)
 create mode 100644 src/lsdb_ml/batch.py
 delete mode 100644 src/lsdb_ml/inputs.py
 delete mode 100644 src/lsdb_ml/spec.py
 create mode 100644 tests/lsdb_ml/test_batch.py
 delete mode 100644 tests/lsdb_ml/test_collate.py
 delete mode 100644 tests/lsdb_ml/test_spec.py

diff --git a/docs/design/mvp-design.md b/docs/design/mvp-design.md
index c8d23a9..cbaced4 100644
--- a/docs/design/mvp-design.md
+++ b/docs/design/mvp-design.md
@@ -1,4 +1,4 @@
-# lsdb-ml MVP: column spec and inference
+# lsdb-ml MVP: catalog inference
 
 Status: implemented on this branch. This document covers only what the MVP ships. The fuller
 design, which also covers training, is on the `design/initial-design` branch; the parts of it
@@ -10,65 +10,89 @@ This MVP runs a PyTorch model over every row of an LSDB catalog on Dask workers.
 a lazy catalog of predictions that has correct `meta` and is aligned row by row with the
 input. Two pieces make that up:
 
-- **Column spec:** maps flat and nested catalog columns to tensors, validated against
-  catalog `meta` before any data is read.
+- **Inputs by column name:** the model's inputs are plain catalog column names. Whether
+  each is flat or nested, and its dtype, are read from catalog `meta`, and everything is
+  checked before any data is read.
 - **Inference:** a Dask `WorkerPlugin` installs the model once per worker, and
   `map_partitions` runs it partition by partition.
 
 Training, the Hyrax backend, and running without a `distributed` client are out of scope.
 `hyrax` remains only an optional extra.
 
+Data preparation (filtering, derived columns, unit conversion, encoding strings) is left to
+ordinary LSDB operations before calling `predict`. lsdb-ml has no column-spec layer of its
+own; one can be added later if a clear need appears.
+
 ## 2. API
 
 ```python
 import torch
 from distributed import Client
-from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested
+from lsdb_ml import ModelDeployment, ModelSpec
 
 client = Client(...)  # a dask.distributed client is required
 
-spec = ColumnSpec({
-    "mag": Flat("mag"),                                              # -> (B,)
-    "flux": Nested("lc", "flux", max_length=32, truncate="tail"),    # -> (B, L), flux_mask, flux_lengths
-})
+class MyModel(torch.nn.Module):
+    def forward(self, batch):          # batch: dict[str, Tensor] keyed by column name
+        mag = batch["mag"]                                       # (B,)
+        flux, mask = batch["lc.flux"], batch["lc._mask"]         # (B, L), (B, L)
+        ...
 
 with ModelDeployment(ModelSpec(MyModel, weights="model.pt")) as deployed:
     scored = deployed.predict(
-        catalog, spec,
+        catalog,
+        inputs=["mag", "lc.flux"],     # flat column, nested sub-column (LSDB dotted name)
         outputs={"score": "float32"},
-        inputs=["mag", "flux", "flux_mask"],   # batch keys passed positionally to the model
+        max_length=32, truncate="tail",
         batch_size=4096,
     )
     result = scored.compute()
 ```
 
-`predict(catalog, model, spec, outputs, ...)` is the one-call form of the same thing. It
-leaves the model installed, because the returned catalog is lazy. `ModelDeployment.close()`
-removes the model from the workers.
+- **One-call form.** `predict(catalog, model, inputs, outputs, ...)` does the same thing.
+  It leaves the model installed, because the returned catalog is lazy.
+  `ModelDeployment.close()` removes the model from the workers.
+- **Models with positional arguments.** These use an adapter,
+  `forward=lambda m, b: m(b["mag"], b["lc.flux"], b["lc._mask"])`.
 
-## 3. Column spec
+## 3. Inputs and batches
 
-- **Fields.** `Flat(columns, dtype=float32)` stacks numeric columns into a `(B,)` or
-  `(B, F)` tensor. `Nested(column, fields, dtype=float32, max_length=None, truncate="raise")`
-  reads numeric sub-columns of a nested column.
-- **Validation before compute.** `ColumnSpec.resolve(catalog)` uses only `catalog.meta`.
+- **Names.** A flat column is named as itself (`"mag"`). A nested sub-column uses LSDB's
+  dotted name (`"lc.flux"`), the same form `catalog["lc.flux"]` accepts. Batch keys are
+  those same names.
+- **Validation before compute.** `plan_inputs(catalog, inputs)` reads only `catalog.meta`.
   Meta is an empty `NestedFrame` with full dtypes, and `NestedDtype.column_dtypes` gives the
-  sub-column types. That replaces the hand-built column preview in hyrax#1062. Resolve checks:
-  - columns and sub-columns exist;
-  - flat vs nested is used correctly;
-  - dtypes are numeric;
-  - output keys don't collide.
-
-  Every problem is reported in one `SpecError`.
-- **Nested representation: zero padding with a mask.** Each batch is padded to its longest
-  row and returns `<key>` `(B, L[, F])`, `<key>_mask` `(B, L)` and `<key>_lengths` `(B,)`.
-  This works directly with attention masks, masked pooling and `pack_padded_sequence`.
-  `torch.nested` is not mature enough to be the default. Other representations can be added
-  later as new `Field` subclasses.
-- **Collation is vectorized.** It is built from `series.nest.len()` (null rows count as 0)
-  and `series.nest.to_flat()`, with no per-row Python.
-- **Missing values raise.** Drop or fill them in the catalog first.
-- **Every batch includes `_index`.** It holds the frame index (`_healpix_29`) as int64.
+  sub-column types. That replaces the hand-built column preview in hyrax#1062. It checks:
+  - each name exists;
+  - a whole nested column isn't named where a sub-column is needed;
+  - dtypes are numeric or bool.
+
+  Every problem is reported in one `InputError`. `predict` then dry-runs the conversion on
+  `catalog.meta` itself, which yields zero-row tensors and still reads no data.
+- **Dtypes are inferred.** Float columns become `float_dtype` (float32 by default), integers
+  int64, and booleans bool.
+- **Nested representation: zero padding with a shared mask.** Each nested sub-column
+  becomes a `(B, L)` tensor, padded to the longest row in the batch. Sub-columns of one
+  nested column share offsets, so they share one `"<column>._mask"` (bool, `(B, L)`) and
+  one `"<column>._lengths"` (int64, `(B,)`). This works directly with attention masks,
+  masked pooling and `pack_padded_sequence`. `torch.nested` is not mature enough to be the
+  default.
+- **Sub-columns stay separate tensors.** They aren't stacked into `(B, L, F)`, to keep
+  things simple. Models can stack them if they want.
+- **Conversion is vectorized.** It is built from `series.nest.len()` (null rows count as 0)
+  and `series.nest.to_flat()`, with no per-row Python. Missing values raise; drop or fill
+  them in the catalog first.
+- **Batch-level padding is always correct, but has two costs.**
+  - *Width follows the longest row.* One very long light curve makes its whole batch that
+    wide, so memory and compute scale with `B × L_max` per batch. `max_length` with
+    `truncate` caps this.
+  - *Width varies from batch to batch.* Models must accept a variable `L`. That's natural
+    for sequence models and pooling, but a model that needs a fixed `L` (e.g. an MLP over a
+    flattened sequence) isn't supported yet. A fixed-width padding option (`pad_to`, or a
+    catalog-level `pad_nested`) is the natural follow-up if that's needed.
+
+  Because padding is masked, predictions don't depend on which other rows share a batch,
+  as long as the model uses the mask.
 
 ## 4. Inference lifecycle
 
@@ -107,16 +131,17 @@ Predictions are never joined on `_healpix_29`, which is not guaranteed unique. T
 All tests run on CPU, take about 16 seconds, and use a synthetic nested catalog and an
 in-process 2-worker × 2-thread cluster.
 
-- **Spec validation:** every error case, including a fixture catalog that raises if any
-  partition is computed.
-- **Collation:** checked against a per-row reference, covering truncation, empty and null
-  nested rows, missing values, and row order.
+- **Input validation:** every error case, including a fixture catalog that raises if any
+  partition is computed, and dtype inference from meta.
+- **Conversion:** checked against a per-row reference, covering the shared mask, truncation,
+  empty and null nested rows, missing values, and row order.
 - **Alignment:** a filter that empties partitions, then a reordered partition subset; and a
   catalog where every row shares one `_healpix_29` value. Both are checked per row and joined
   on object id. Injecting a reversed batch order makes these tests fail.
 - **Load once per worker:** the model is built exactly once per worker across many
   partitions and repeated predictions.
 - **Meta:** dtypes of the lazy result match the computed result.
+- **`forward=` adapter:** a model taking positional tensors works through it.
 - **Failure modes:** wrong output row count, closed deployment, no client, bad arguments.
 
 ## 7. Deferred (in the full design)
@@ -124,10 +149,11 @@ in-process 2-worker × 2-thread cluster.
 | Item | Why deferred |
 |---|---|
 | Training: `CatalogIterableDataset` (partition sharding across DataLoader workers and DDP) and `TrainSession` | Larger surface; depends on upstream `CatalogStream` changes. Planned as two follow-up PRs. |
-| `transform=` hooks, `pad_to`/`pad_value`, missing-value filling | Convenience; easy to add back when needed. |
+| A column-spec layer (`ColumnSpec`, transforms, column grouping, missing-value filling) | Data preparation belongs to LSDB operations; add a middle layer only if a clear need appears. |
+| Fixed-width padding (`pad_to` or catalog-level `pad_nested`), stacking sub-columns into `(B, L, F)` | Needed only for fixed-`L` models or as an optimization; batch-level padding covers correctness. |
 | Dict-valued and list-column model outputs | Not needed for the first use cases. |
 | Running without a `distributed` client | A second code path; requiring a `Client` is a reasonable MVP constraint. |
-| Hyrax backend | The seams already exist: the model factory in `ModelSpec` and `inputs=` (a key list or callable). |
+| Hyrax backend | The seams already exist: the model factory in `ModelSpec` and the `forward=` adapter. |
 | Automatic output meta from a dummy forward pass | Fragile for models that can't take an empty or synthetic batch. |
 
 ## 8. CI note
diff --git a/docs/examples/predict_on_catalog.py b/docs/examples/predict_on_catalog.py
index b68b051..b29f952 100644
--- a/docs/examples/predict_on_catalog.py
+++ b/docs/examples/predict_on_catalog.py
@@ -17,7 +17,7 @@
 from distributed import Client, LocalCluster
 from torch import nn
 
-from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested
+from lsdb_ml import ModelDeployment, ModelSpec
 
 
 def make_catalog(n_objects=5000, seed=0):
@@ -46,8 +46,10 @@ def __init__(self):
         super().__init__()
         self.head = nn.Linear(2, 2)
 
-    def forward(self, mag, flux, mask):
+    def forward(self, batch):
         """Return two scores per object from the magnitude and the mean light-curve flux."""
+        # Batches are keyed by column name; nested sub-columns share a "<column>._mask".
+        mag, flux, mask = batch["mag"], batch["lc.flux"], batch["lc._mask"]
         mean_flux = (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
         return self.head(torch.stack([mag, mean_flux], dim=1))
 
@@ -58,15 +60,14 @@ def main(weights_path):
     torch.save(Scorer().state_dict(), weights_path)  # stands in for a trained model
 
     catalog = make_catalog().query("mag < 21")  # any lazy lsdb pipeline works here
-    spec = ColumnSpec({"mag": Flat("mag"), "flux": Nested("lc", "flux", max_length=32, truncate="tail")})
-
-    # Workers build Scorer() and load the weights themselves; only the spec is shipped.
+    # Workers build Scorer() and load the weights themselves; only the factory and path are shipped.
     with ModelDeployment(ModelSpec(Scorer, weights=weights_path, device="cpu")) as deployed:
         scored = deployed.predict(
             catalog,
-            spec,
+            inputs=["mag", "lc.flux"],  # flat vs nested is read from the catalog's meta
             outputs={"score_a": "float32", "score_b": "float32"},
-            inputs=["mag", "flux", "flux_mask"],
+            max_length=32,  # keep at most the last 32 observations per light curve
+            truncate="tail",
             batch_size=1024,
         )
         print(scored.dtypes[["object_id", "mag", "score_a", "score_b"]])  # meta, nothing computed yet
diff --git a/src/lsdb_ml/__init__.py b/src/lsdb_ml/__init__.py
index 888e182..2232ac5 100644
--- a/src/lsdb_ml/__init__.py
+++ b/src/lsdb_ml/__init__.py
@@ -1,21 +1,17 @@
 from ._version import __version__
+from .batch import InputError, InputPlan, plan_inputs
 from .device import move_batch, resolve_device
 from .example_module import greetings, meaning
 from .inference import ModelDeployment, ModelSpec, OutputSpec, predict
-from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo
 
 __all__ = [
-    "ColumnSpec",
-    "Field",
-    "Flat",
-    "Nested",
-    "ResolvedSpec",
-    "SpecError",
-    "TensorInfo",
+    "predict",
     "ModelDeployment",
     "ModelSpec",
     "OutputSpec",
-    "predict",
+    "InputError",
+    "InputPlan",
+    "plan_inputs",
     "move_batch",
     "resolve_device",
     "greetings",
diff --git a/src/lsdb_ml/batch.py b/src/lsdb_ml/batch.py
new file mode 100644
index 0000000..de0d392
--- /dev/null
+++ b/src/lsdb_ml/batch.py
@@ -0,0 +1,264 @@
+"""Turning catalog columns into batches of tensors.
+
+Inputs are named with plain column names. Flat columns use their name (``"mag"``), and
+nested sub-columns use LSDB's dotted names (``"lc.flux"``). Whether a column is flat or
+nested, and its dtype, are read from the catalog's ``meta``, so inputs can be checked
+before any data is computed.
+
+A batch is a ``dict[str, torch.Tensor]`` keyed by those same names:
+
+- a flat column becomes a ``(B,)`` tensor;
+- a nested sub-column becomes a ``(B, L)`` tensor, zero-padded to the longest row in the
+  batch (after optional truncation);
+- each nested column with selected sub-columns also gets ``"<column>._mask"`` (bool,
+  ``(B, L)``, True where a value is real) and ``"<column>._lengths"`` (int64, ``(B,)``).
+  Sub-columns of one nested column share offsets, so they share one mask.
+
+Floating-point columns are cast to ``float_dtype`` (float32 by default). Integer columns
+become int64 and boolean columns stay bool.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import torch
+from nested_pandas.series.dtype import NestedDtype
+
+from lsdb_ml.collate import TruncatePolicy, frame_to_array, pad_ragged
+
+MASK_SUFFIX = "._mask"
+LENGTHS_SUFFIX = "._lengths"
+
+
+class InputError(ValueError):
+    """Raised when requested input columns do not match a catalog's schema."""
+
+
+def _kind(dtype) -> str | None:
+    """Classify a pandas/numpy/pyarrow dtype as "float", "int", "bool", or None (unsupported)."""
+    if isinstance(dtype, NestedDtype):
+        return None
+    if isinstance(dtype, pd.ArrowDtype):
+        dtype = dtype.pyarrow_dtype
+    if isinstance(dtype, pa.DataType):
+        if pa.types.is_floating(dtype):
+            return "float"
+        if pa.types.is_integer(dtype):
+            return "int"
+        if pa.types.is_boolean(dtype):
+            return "bool"
+        return None
+    if pd.api.types.is_bool_dtype(dtype):
+        return "bool"
+    if pd.api.types.is_float_dtype(dtype):
+        return "float"
+    if pd.api.types.is_integer_dtype(dtype):
+        return "int"
+    return None
+
+
+def _torch_dtype(kind: str, float_dtype: torch.dtype) -> torch.dtype:
+    return {"float": float_dtype, "int": torch.int64, "bool": torch.bool}[kind]
+
+
+@dataclass(frozen=True)
+class _Column:
+    name: str
+    kind: str
+
+
+class InputPlan:
+    """Validated inputs for a catalog: which columns become which tensors.
+
+    Create it with :func:`plan_inputs`. It is picklable, so it can be shipped to Dask
+    tasks.
+    """
+
+    def __init__(
+        self,
+        flat: list[_Column],
+        nested: dict[str, list[_Column]],
+        float_dtype: torch.dtype,
+        max_length: int | None,
+        truncate: TruncatePolicy,
+    ):
+        self.flat = flat
+        self.nested = nested
+        self.float_dtype = float_dtype
+        self.max_length = max_length
+        self.truncate = truncate
+
+    def describe(self) -> dict[str, torch.dtype]:
+        """Return the dtype of every tensor in a batch, keyed by batch key."""
+        info = {c.name: _torch_dtype(c.kind, self.float_dtype) for c in self.flat}
+        for parent, subs in self.nested.items():
+            for sub in subs:
+                info[f"{parent}.{sub.name}"] = _torch_dtype(sub.kind, self.float_dtype)
+            info[parent + MASK_SUFFIX] = torch.bool
+            info[parent + LENGTHS_SUFFIX] = torch.int64
+        return info
+
+    def to_tensors(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
+        """Convert the rows of ``frame`` to a batch of tensors.
+
+        Row ``i`` of every tensor comes from row ``i`` of ``frame``.
+
+        Parameters
+        ----------
+        frame : pd.DataFrame
+            Rows with the planned columns, e.g. a slice of a catalog partition. An empty
+            frame, such as ``catalog.meta``, gives zero-row tensors.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            One entry per key in :meth:`describe`.
+
+        Raises
+        ------
+        ValueError
+            If a value is missing, or a nested row is longer than ``max_length`` with
+            ``truncate="raise"``.
+        """
+        batch: dict[str, torch.Tensor] = {}
+        for column in self.flat:
+            dtype = _numpy_dtype(column.kind, self.float_dtype)
+            batch[column.name] = torch.from_numpy(
+                frame_to_array(frame[[column.name]], dtype, repr(column.name))[:, 0].copy()
+            )
+        for parent, subs in self.nested.items():
+            series = frame[parent]
+            # nest.len() counts null rows as 0; to_flat() returns every row's values
+            # concatenated in row order. Both stay consistent across iloc slices.
+            lengths = np.asarray(series.nest.len(), dtype=np.int64)
+            flat_values = series.nest.to_flat([sub.name for sub in subs])
+            for sub in subs:
+                key = f"{parent}.{sub.name}"
+                values = frame_to_array(
+                    flat_values[[sub.name]], _numpy_dtype(sub.kind, self.float_dtype), repr(key)
+                )
+                padded, mask, kept = pad_ragged(
+                    values, lengths, max_length=self.max_length, truncate=self.truncate, what=repr(key)
+                )
+                batch[key] = torch.from_numpy(np.ascontiguousarray(padded[..., 0]))
+            batch[parent + MASK_SUFFIX] = torch.from_numpy(mask)
+            batch[parent + LENGTHS_SUFFIX] = torch.from_numpy(kept)
+        return batch
+
+
+def _numpy_dtype(kind: str, float_dtype: torch.dtype) -> np.dtype:
+    return torch.empty(0, dtype=_torch_dtype(kind, float_dtype)).numpy().dtype
+
+
+def plan_inputs(
+    schema,
+    inputs: str | Sequence[str],
+    *,
+    float_dtype: torch.dtype = torch.float32,
+    max_length: int | None = None,
+    truncate: TruncatePolicy = "raise",
+) -> InputPlan:
+    """Check input column names against a catalog's schema, without computing data.
+
+    Parameters
+    ----------
+    schema : lsdb.Catalog or pd.DataFrame
+        A catalog (its ``meta`` is used), or a frame with the same columns and dtypes.
+    inputs : str or sequence of str
+        Column names: ``"mag"`` for a flat column, ``"lc.flux"`` for a nested sub-column.
+    float_dtype : torch.dtype, default torch.float32
+        Dtype for floating-point columns.
+    max_length : int, optional
+        Longest nested sequence to keep; see ``truncate``.
+    truncate : {"raise", "head", "tail"}, default "raise"
+        For nested rows longer than ``max_length``: raise, keep the first values, or keep
+        the last values.
+
+    Returns
+    -------
+    InputPlan
+        The validated plan.
+
+    Raises
+    ------
+    InputError
+        Listing every problem found.
+
+    Examples
+    --------
+    >>> import nested_pandas as npd
+    >>> nf = npd.NestedFrame({"a": [1.0, 2.0]}, index=[10, 11])
+    >>> nf = nf.join_nested(npd.NestedFrame({"f": [1.0, 2.0, 3.0]}, index=[10, 10, 11]), "lc")
+    >>> batch = plan_inputs(nf, ["a", "lc.f"]).to_tensors(nf)
+    >>> batch["lc.f"]
+    tensor([[1., 2.],
+            [3., 0.]])
+    >>> batch["lc._mask"]
+    tensor([[ True,  True],
+            [ True, False]])
+    """
+    meta = schema if isinstance(schema, pd.DataFrame) else getattr(schema, "meta", None)
+    if not isinstance(meta, pd.DataFrame):
+        raise TypeError(f"Expected an lsdb catalog or a DataFrame, got {type(schema).__name__}.")
+    names = [inputs] if isinstance(inputs, str) else list(inputs)
+
+    problems: list[str] = []
+    if not names:
+        problems.append("No input columns were given.")
+    if truncate not in ("raise", "head", "tail"):
+        problems.append(f"truncate must be 'raise', 'head' or 'tail', got {truncate!r}.")
+    if max_length is not None and (
+        isinstance(max_length, bool) or not isinstance(max_length, int) or max_length < 1
+    ):
+        problems.append(f"max_length must be a positive integer, got {max_length!r}.")
+    if float_dtype not in (torch.float16, torch.float32, torch.float64):
+        problems.append(f"float_dtype must be torch.float16, float32 or float64, got {float_dtype}.")
+
+    flat: list[_Column] = []
+    nested: dict[str, list[_Column]] = {}
+    for name in dict.fromkeys(names):  # drop duplicates, keep order
+        if name in meta.columns:
+            dtype = meta.dtypes[name]
+            if isinstance(dtype, NestedDtype):
+                problems.append(
+                    f"'{name}' is a nested column; name its sub-columns instead, e.g. "
+                    f"{[f'{name}.{sub}' for sub in list(dtype.column_dtypes)[:2]]}."
+                )
+            elif _kind(dtype) is None:
+                problems.append(f"'{name}' has non-numeric dtype {dtype}.")
+            else:
+                flat.append(_Column(name, _kind(dtype)))
+            continue
+        parent, _, sub = name.partition(".")
+        if not sub or parent not in meta.columns:
+            problems.append(f"'{name}' is not a column of the catalog.")
+        elif not isinstance(meta.dtypes[parent], NestedDtype):
+            problems.append(f"'{name}': '{parent}' is not a nested column.")
+        elif sub not in meta.dtypes[parent].column_dtypes:
+            available = sorted(meta.dtypes[parent].column_dtypes)
+            problems.append(
+                f"'{name}': nested column '{parent}' has no sub-column '{sub}' (available: {available})."
+            )
+        elif _kind(meta.dtypes[parent].column_dtypes[sub]) is None:
+            problems.append(f"'{name}' has non-numeric type {meta.dtypes[parent].column_dtypes[sub]}.")
+        else:
+            nested.setdefault(parent, []).append(_Column(sub, _kind(meta.dtypes[parent].column_dtypes[sub])))
+
+    for parent, subs in nested.items():
+        for suffix in (MASK_SUFFIX, LENGTHS_SUFFIX):
+            if any(f"{parent}.{sub.name}" == parent + suffix for sub in subs):
+                problems.append(
+                    f"'{parent + suffix}' is reserved for the generated mask/lengths of '{parent}'."
+                )
+
+    if problems:
+        hint = ""
+        if any("is not a column" in p for p in problems) and getattr(schema, "all_columns", None) is not None:
+            hint = f"\nThe catalog was loaded with columns {list(meta.columns)}; others may exist on disk."
+        raise InputError("Input columns do not match the catalog:\n  - " + "\n  - ".join(problems) + hint)
+    return InputPlan(flat, nested, float_dtype, max_length, truncate)
diff --git a/src/lsdb_ml/collate.py b/src/lsdb_ml/collate.py
index 2de8aea..1efa96f 100644
--- a/src/lsdb_ml/collate.py
+++ b/src/lsdb_ml/collate.py
@@ -1,7 +1,7 @@
 """Vectorized conversion of (nested) dataframe columns into padded numpy arrays.
 
-These helpers do the numeric work behind :class:`lsdb_ml.Flat` and :class:`lsdb_ml.Nested`.
-Nothing here loops over rows in Python.
+These helpers do the numeric work behind :func:`lsdb_ml.batch.plan_inputs`. Nothing here
+loops over rows in Python.
 """
 
 from __future__ import annotations
@@ -10,40 +10,10 @@
 
 import numpy as np
 import pandas as pd
-import torch
 
 TruncatePolicy = Literal["raise", "head", "tail"]
 
 
-def numpy_dtype(dtype: torch.dtype) -> np.dtype:
-    """Return the numpy dtype that corresponds to a torch dtype.
-
-    Parameters
-    ----------
-    dtype : torch.dtype
-        A torch dtype with a numpy equivalent.
-
-    Returns
-    -------
-    np.dtype
-        The matching numpy dtype.
-
-    Raises
-    ------
-    TypeError
-        If the torch dtype has no numpy equivalent (e.g. ``torch.bfloat16``).
-
-    Examples
-    --------
-    >>> numpy_dtype(torch.float32)
-    dtype('float32')
-    """
-    try:
-        return torch.empty(0, dtype=dtype).numpy().dtype
-    except TypeError as err:
-        raise TypeError(f"torch dtype {dtype} has no numpy equivalent and cannot be collated.") from err
-
-
 def frame_to_array(frame: pd.DataFrame, dtype: np.dtype, what: str) -> np.ndarray:
     """Convert the columns of ``frame`` into a writable ``(N, F)`` array.
 
diff --git a/src/lsdb_ml/inference/predict.py b/src/lsdb_ml/inference/predict.py
index 1e917fd..de00d79 100644
--- a/src/lsdb_ml/inference/predict.py
+++ b/src/lsdb_ml/inference/predict.py
@@ -3,18 +3,18 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 
 import numpy as np
 import pandas as pd
 import torch
 from torch import nn
 
+from lsdb_ml.batch import InputPlan, plan_inputs
+from lsdb_ml.collate import TruncatePolicy
 from lsdb_ml.device import move_batch
 from lsdb_ml.inference.outputs import OutputSpec
 from lsdb_ml.inference.plugin import ModelPlugin, ModelSpec, get_model
-from lsdb_ml.inputs import InputsLike, model_args
-from lsdb_ml.spec import ColumnSpec, ResolvedSpec
 
 logger = logging.getLogger(__name__)
 
@@ -22,6 +22,9 @@
 # model to the same client does not rebuild it on every worker.
 _REGISTERED: set[tuple[str, str]] = set()
 
+#: ``forward(model, batch) -> Tensor``: how to call the model on a batch dict.
+Forward = Callable[[nn.Module, dict[str, torch.Tensor]], torch.Tensor]
+
 
 def _active_client():
     from distributed import get_client
@@ -50,7 +53,7 @@ class ModelDeployment:
     Use it as a context manager to remove the model from the workers afterwards::
 
         with ModelDeployment(ModelSpec(make_model, weights="model.pt")) as deployed:
-            result = deployed.predict(catalog, spec, inputs="x", outputs={"z": "float32"}).compute()
+            result = deployed.predict(catalog, ["mag", "lc.flux"], {"z": "float32"}).compute()
 
     Parameters
     ----------
@@ -79,12 +82,15 @@ def __init__(self, model: ModelSpec | nn.Module, client=None):
     def predict(
         self,
         catalog,
-        spec: ColumnSpec | ResolvedSpec,
+        inputs: str | Sequence[str],
         outputs: OutputSpec | Mapping[str, object] | Sequence[str],
         *,
-        inputs: InputsLike | None = None,
+        forward: Forward | None = None,
         batch_size: int = 4096,
         append: bool = True,
+        float_dtype: torch.dtype = torch.float32,
+        max_length: int | None = None,
+        truncate: TruncatePolicy = "raise",
     ):
         """Return a lazy catalog with the model's predictions for every row.
 
@@ -92,32 +98,27 @@ def predict(
         """
         if self._closed:
             raise RuntimeError("This deployment is closed.")
-        resolved = spec if isinstance(spec, ResolvedSpec) else spec.resolve(catalog)
+        plan = plan_inputs(catalog, inputs, float_dtype=float_dtype, max_length=max_length, truncate=truncate)
+        plan.to_tensors(catalog.meta)  # dry run on the empty meta frame: no data is read
         if isinstance(batch_size, bool) or int(batch_size) < 1:
             raise ValueError(f"batch_size must be a positive integer, got {batch_size!r}.")
+        if forward is not None and not callable(forward):
+            raise TypeError(f"forward must be callable(model, batch), got {type(forward).__name__}.")
         outputs = outputs if isinstance(outputs, OutputSpec) else OutputSpec(outputs)
-        if inputs is None:
-            if len(resolved.fields) != 1:
-                raise ValueError(
-                    f"The spec has several fields {list(resolved.fields)}; pass inputs= to say which "
-                    "batch keys the model takes."
-                )
-            inputs = next(iter(resolved.fields))
         if append:
             clashes = sorted(set(outputs.columns) & set(catalog.columns))
             if clashes:
                 raise ValueError(f"Output columns {clashes} already exist in the catalog; rename them.")
 
-        meta = outputs.meta(catalog.meta, append)
         return catalog.map_partitions(
             _predict_partition,
             model_name=self.name,
-            resolved=resolved,
-            inputs=inputs,
+            plan=plan,
+            forward=forward,
             outputs=outputs,
             batch_size=int(batch_size),
             append=append,
-            meta=meta,
+            meta=outputs.meta(catalog.meta, append),
         )
 
     def close(self) -> None:
@@ -138,12 +139,15 @@ def __exit__(self, *exc_info) -> None:
 def predict(
     catalog,
     model: ModelSpec | nn.Module,
-    spec: ColumnSpec | ResolvedSpec,
+    inputs: str | Sequence[str],
     outputs: OutputSpec | Mapping[str, object] | Sequence[str],
     *,
-    inputs: InputsLike | None = None,
+    forward: Forward | None = None,
     batch_size: int = 4096,
     append: bool = True,
+    float_dtype: torch.dtype = torch.float32,
+    max_length: int | None = None,
+    truncate: TruncatePolicy = "raise",
     client=None,
 ):
     """Run a model over every row of a catalog, returning a lazy catalog of predictions.
@@ -163,19 +167,29 @@ def predict(
         The input catalog.
     model : ModelSpec or nn.Module
         The model, or how to build it on each worker.
-    spec : ColumnSpec or ResolvedSpec
-        Which columns become which tensors. Validated against ``catalog.meta`` now,
-        before any data is read.
+    inputs : str or sequence of str
+        Input columns: ``"mag"`` for a flat column, ``"lc.flux"`` for a nested sub-column.
+        Checked against ``catalog.meta`` now, before any data is read. See
+        :mod:`lsdb_ml.batch` for the tensors each produces.
     outputs : OutputSpec, mapping or sequence
         Output column names and dtypes, e.g. ``{"z_phot": "float32"}``.
-    inputs : str, sequence of str or callable, optional
-        Batch key(s) passed positionally to the model, or ``callable(batch)``. Defaults to
-        the spec's only field.
+    forward : callable, optional
+        ``forward(model, batch) -> Tensor``, for models that do not take the batch dict
+        directly, e.g. ``lambda m, b: m(b["mag"], b["lc.flux"], b["lc._mask"])``. By
+        default the model is called as ``model(batch)``.
     batch_size : int, default 4096
         Maximum rows per forward pass.
     append : bool, default True
         Keep the input columns. With False, the result holds only the predictions (on the
         same index).
+    float_dtype : torch.dtype, default torch.float32
+        Dtype for floating-point input columns.
+    max_length : int, optional
+        Longest nested sequence to keep; see ``truncate``. Nested inputs are padded to the
+        longest row in each batch, so one very long row makes the whole batch that wide.
+    truncate : {"raise", "head", "tail"}, default "raise"
+        For nested rows longer than ``max_length``: raise, keep the first values, or keep
+        the last values.
     client : distributed.Client, optional
         Client to deploy to. Defaults to the active client; one is required.
 
@@ -185,15 +199,25 @@ def predict(
         Lazy catalog whose ``meta`` includes the output columns.
     """
     deployment = ModelDeployment(model, client=client)
-    return deployment.predict(catalog, spec, outputs, inputs=inputs, batch_size=batch_size, append=append)
+    return deployment.predict(
+        catalog,
+        inputs,
+        outputs,
+        forward=forward,
+        batch_size=batch_size,
+        append=append,
+        float_dtype=float_dtype,
+        max_length=max_length,
+        truncate=truncate,
+    )
 
 
 def _predict_partition(
     partition: pd.DataFrame,
     *,
     model_name: str,
-    resolved: ResolvedSpec,
-    inputs: InputsLike,
+    plan: InputPlan,
+    forward: Forward | None,
     outputs: OutputSpec,
     batch_size: int,
     append: bool,
@@ -206,7 +230,8 @@ def _predict_partition(
     pieces: list[dict[str, np.ndarray]] = []
     with torch.inference_mode():
         for rows in _iter_row_slices(partition, batch_size):
-            batch = move_batch(resolved.collate(rows), device)
-            pieces.append(outputs.to_columns(model(*model_args(batch, inputs)), len(rows)))
+            batch = move_batch(plan.to_tensors(rows), device)
+            output = model(batch) if forward is None else forward(model, batch)
+            pieces.append(outputs.to_columns(output, len(rows)))
     columns = {name: np.concatenate([piece[name] for piece in pieces]) for name in outputs.columns}
     return outputs.build_frame(partition, columns, append)
diff --git a/src/lsdb_ml/inputs.py b/src/lsdb_ml/inputs.py
deleted file mode 100644
index 6e18329..0000000
--- a/src/lsdb_ml/inputs.py
+++ /dev/null
@@ -1,41 +0,0 @@
-"""Selecting model inputs from a collated batch."""
-
-from __future__ import annotations
-
-from collections.abc import Callable, Mapping, Sequence
-from typing import Any
-
-import torch
-
-#: How to build the model's positional arguments from a batch: one key, a list of keys,
-#: or a callable ``batch -> argument`` (a tuple return value is unpacked into arguments).
-InputsLike = str | Sequence[str] | Callable[[Mapping[str, torch.Tensor]], Any]
-
-
-def model_args(batch: Mapping[str, torch.Tensor], inputs: InputsLike) -> tuple:
-    """Return the positional arguments to call a model with.
-
-    Parameters
-    ----------
-    batch : Mapping[str, torch.Tensor]
-        A collated batch.
-    inputs : str, sequence of str, or callable
-        A key, keys (passed positionally in order), or ``callable(batch)``.
-
-    Returns
-    -------
-    tuple
-        Positional arguments.
-
-    Examples
-    --------
-    >>> batch = {"x": 1, "y": 2}
-    >>> model_args(batch, "x"), model_args(batch, ["y", "x"]), model_args(batch, lambda b: (b["x"], 5))
-    ((1,), (2, 1), (1, 5))
-    """
-    if isinstance(inputs, str):
-        return (batch[inputs],)
-    if callable(inputs):
-        result = inputs(batch)
-        return result if isinstance(result, tuple) else (result,)
-    return tuple(batch[key] for key in inputs)
diff --git a/src/lsdb_ml/spec.py b/src/lsdb_ml/spec.py
deleted file mode 100644
index f215d3e..0000000
--- a/src/lsdb_ml/spec.py
+++ /dev/null
@@ -1,423 +0,0 @@
-"""Column specifications: map flat and nested catalog columns to tensors.
-
-A :class:`ColumnSpec` names the tensors a model needs and where each comes from. It is
-checked against a catalog's ``meta`` (an empty ``NestedFrame`` that carries every dtype,
-including nested sub-column dtypes) before any data is computed, and the resulting
-:class:`ResolvedSpec` collates dataframe rows into a ``dict[str, torch.Tensor]``.
-
->>> import lsdb_ml
->>> spec = lsdb_ml.ColumnSpec({"x": lsdb_ml.Flat(["ra", "dec"]), "lc": lsdb_ml.Nested("lc", ["flux"])})
-"""
-
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
-
-import numpy as np
-import pandas as pd
-import pyarrow as pa
-import torch
-from nested_pandas.series.dtype import NestedDtype
-
-from lsdb_ml.collate import TruncatePolicy, frame_to_array, numpy_dtype, pad_ragged
-
-
-class SpecError(ValueError):
-    """Raised when a column specification does not match a catalog's schema."""
-
-
-@dataclass(frozen=True)
-class TensorInfo:
-    """Description of one tensor a resolved spec produces.
-
-    Attributes
-    ----------
-    dtype : torch.dtype
-        The tensor dtype.
-    shape : tuple
-        The shape; ``None`` marks a dimension that varies by batch (always the first one).
-    """
-
-    dtype: torch.dtype
-    shape: tuple[int | None, ...]
-
-
-def _as_list(columns: str | Sequence[str]) -> list[str]:
-    return [columns] if isinstance(columns, str) else list(columns)
-
-
-def _is_numeric(dtype) -> bool:
-    """True for pandas/numpy/pyarrow integer, float, and bool dtypes."""
-    if isinstance(dtype, pa.DataType):
-        return pa.types.is_integer(dtype) or pa.types.is_floating(dtype) or pa.types.is_boolean(dtype)
-    if isinstance(dtype, pd.ArrowDtype):
-        return _is_numeric(dtype.pyarrow_dtype)
-    if isinstance(dtype, NestedDtype):
-        return False
-    try:
-        return pd.api.types.is_numeric_dtype(dtype) or pd.api.types.is_bool_dtype(dtype)
-    except TypeError:
-        return False
-
-
-class Field(ABC):
-    """Base class for one entry of a :class:`ColumnSpec`.
-
-    Subclasses implement :meth:`resolve`, which validates the field against catalog meta
-    and returns a :class:`ResolvedField` that does the actual collation. New column
-    representations (e.g. ragged tensors) are added as new ``Field`` types.
-    """
-
-    @abstractmethod
-    def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None:
-        """Validate against ``meta`` and build the collator.
-
-        Parameters
-        ----------
-        key : str
-            The output key this field was registered under.
-        meta : pd.DataFrame
-            The catalog meta (an empty frame with full dtypes).
-        problems : list[str]
-            Append a message for every problem found, instead of raising, so that all
-            problems in a spec are reported together.
-
-        Returns
-        -------
-        ResolvedField or None
-            None if any problem was found.
-        """
-
-
-class ResolvedField(ABC):
-    """A field validated against a schema, ready to collate frames."""
-
-    #: Top-level columns this field reads.
-    columns: list[str]
-
-    @abstractmethod
-    def describe(self) -> dict[str, TensorInfo]:
-        """Return the tensors this field produces, keyed by output key."""
-
-    @abstractmethod
-    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
-        """Convert the rows of ``frame`` to tensors."""
-
-
-@dataclass
-class Flat(Field):
-    """One or more flat (non-nested) numeric columns, stacked into a single tensor.
-
-    Parameters
-    ----------
-    columns : str or sequence of str
-        A single column name produces a ``(B,)`` tensor; a list produces ``(B, F)``.
-    dtype : torch.dtype, default torch.float32
-        Output dtype.
-
-    Notes
-    -----
-    Missing values raise an error during collation; drop or fill them in the catalog first.
-    """
-
-    columns: str | Sequence[str]
-    dtype: torch.dtype = torch.float32
-
-    def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None:
-        """Validate the columns against ``meta``; see :meth:`Field.resolve`."""
-        columns = _as_list(self.columns)
-        start = len(problems)
-        if not columns:
-            problems.append(f"'{key}': Flat needs at least one column.")
-        for column in columns:
-            if column not in meta.columns:
-                problems.append(f"'{key}': column '{column}' is not in the catalog.")
-            elif isinstance(meta.dtypes[column], NestedDtype):
-                problems.append(f"'{key}': column '{column}' is nested; use Nested('{column}', [...]).")
-            elif not _is_numeric(meta.dtypes[column]):
-                problems.append(f"'{key}': column '{column}' has non-numeric dtype {meta.dtypes[column]}.")
-        try:
-            np_dtype = numpy_dtype(self.dtype)
-        except TypeError as err:
-            problems.append(f"'{key}': {err}")
-        if len(problems) > start:
-            return None
-        return _ResolvedFlat(key, columns, isinstance(self.columns, str), self, np_dtype)
-
-
-class _ResolvedFlat(ResolvedField):
-    def __init__(self, key: str, columns: list[str], scalar: bool, spec: Flat, np_dtype: np.dtype):
-        self.key = key
-        self.columns = columns
-        self.scalar = scalar
-        self.spec = spec
-        self.np_dtype = np_dtype
-
-    def describe(self) -> dict[str, TensorInfo]:
-        shape = (None,) if self.scalar else (None, len(self.columns))
-        return {self.key: TensorInfo(self.spec.dtype, shape)}
-
-    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
-        array = frame_to_array(frame[self.columns], self.np_dtype, f"'{self.key}' (columns {self.columns})")
-        if self.scalar:
-            array = array[:, 0].copy()
-        return {self.key: torch.from_numpy(array)}
-
-
-@dataclass
-class Nested(Field):
-    """Numeric sub-columns of a nested column, zero-padded to a rectangular tensor with a mask.
-
-    For output key ``k`` this produces:
-
-    - ``k``: values, ``(B, L)`` for a single sub-column or ``(B, L, F)`` for a list;
-    - ``k_mask``: bool ``(B, L)``, True where a value is real rather than padding;
-    - ``k_lengths``: int64 ``(B,)``, the number of real values per row.
-
-    ``L`` is the longest row in the batch, after truncation. Rows whose nested value is
-    null or empty have length 0. Padding is zero, which carries no meaning; always use
-    the mask.
-
-    Parameters
-    ----------
-    column : str
-        The nested column.
-    fields : str or sequence of str
-        Sub-column name(s) to read.
-    dtype : torch.dtype, default torch.float32
-        Output dtype of the values.
-    max_length : int, optional
-        Longest sequence to keep; see ``truncate``.
-    truncate : {"raise", "head", "tail"}, default "raise"
-        What to do with rows longer than ``max_length``: raise, keep the first values, or
-        keep the last values.
-
-    Notes
-    -----
-    Missing values inside a list raise an error during collation.
-    """
-
-    column: str
-    fields: str | Sequence[str]
-    dtype: torch.dtype = torch.float32
-    max_length: int | None = None
-    truncate: TruncatePolicy = "raise"
-
-    def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None:
-        """Validate the nested column and sub-columns against ``meta``; see :meth:`Field.resolve`."""
-        subcolumns = _as_list(self.fields)
-        start = len(problems)
-        if not subcolumns:
-            problems.append(f"'{key}': Nested needs at least one sub-column.")
-        if self.truncate not in ("raise", "head", "tail"):
-            problems.append(f"'{key}': truncate must be 'raise', 'head' or 'tail', got {self.truncate!r}.")
-        limit = self.max_length
-        if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 1):
-            problems.append(f"'{key}': max_length must be a positive integer, got {limit!r}.")
-        if self.column not in meta.columns:
-            problems.append(f"'{key}': column '{self.column}' is not in the catalog.")
-        elif not isinstance(meta.dtypes[self.column], NestedDtype):
-            problems.append(f"'{key}': column '{self.column}' is not nested; use Flat('{self.column}').")
-        else:
-            available = meta.dtypes[self.column].column_dtypes
-            for sub in subcolumns:
-                if sub not in available:
-                    problems.append(
-                        f"'{key}': nested column '{self.column}' has no sub-column '{sub}' "
-                        f"(available: {sorted(available)})."
-                    )
-                elif not _is_numeric(available[sub]):
-                    problems.append(
-                        f"'{key}': sub-column '{self.column}.{sub}' has non-numeric type {available[sub]}."
-                    )
-        try:
-            np_dtype = numpy_dtype(self.dtype)
-        except TypeError as err:
-            problems.append(f"'{key}': {err}")
-        if len(problems) > start:
-            return None
-        return _ResolvedNested(key, subcolumns, isinstance(self.fields, str), self, np_dtype)
-
-
-class _ResolvedNested(ResolvedField):
-    def __init__(self, key: str, subcolumns: list[str], scalar: bool, spec: Nested, np_dtype: np.dtype):
-        self.key = key
-        self.columns = [spec.column]
-        self.subcolumns = subcolumns
-        self.scalar = scalar
-        self.spec = spec
-        self.np_dtype = np_dtype
-
-    def describe(self) -> dict[str, TensorInfo]:
-        values_shape = (None, None) if self.scalar else (None, None, len(self.subcolumns))
-        return {
-            self.key: TensorInfo(self.spec.dtype, values_shape),
-            f"{self.key}_mask": TensorInfo(torch.bool, (None, None)),
-            f"{self.key}_lengths": TensorInfo(torch.int64, (None,)),
-        }
-
-    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
-        what = f"'{self.key}' (nested column '{self.spec.column}', fields {self.subcolumns})"
-        series = frame[self.spec.column]
-        # nest.len() counts null rows as 0, and to_flat() returns the values of all rows
-        # concatenated in row order; both are consistent across iloc slices and concat.
-        lengths = np.asarray(series.nest.len(), dtype=np.int64)
-        values = frame_to_array(series.nest.to_flat(self.subcolumns), self.np_dtype, what)
-        padded, mask, kept = pad_ragged(
-            values, lengths, max_length=self.spec.max_length, truncate=self.spec.truncate, what=what
-        )
-        if self.scalar:
-            padded = padded[..., 0]
-        return {
-            self.key: torch.from_numpy(np.ascontiguousarray(padded)),
-            f"{self.key}_mask": torch.from_numpy(mask),
-            f"{self.key}_lengths": torch.from_numpy(kept),
-        }
-
-
-@dataclass
-class ColumnSpec:
-    """Mapping from output keys to catalog columns.
-
-    Parameters
-    ----------
-    fields : Mapping[str, Field]
-        Output key to field, e.g. ``{"x": Flat(["ra", "dec"]), "lc": Nested("lc", "flux")}``.
-    index_key : str or None, default "_index"
-        Every batch also includes the frame index (``_healpix_29`` for LSDB catalogs) as an
-        int64 ``(B,)`` tensor under this key, so batches can be traced back to rows. None
-        disables it.
-
-    Examples
-    --------
-    >>> import nested_pandas as npd
-    >>> nf = npd.NestedFrame({"a": [1.0, 2.0]}, index=[10, 11])
-    >>> nf = nf.join_nested(npd.NestedFrame({"f": [1.0, 2.0, 3.0]}, index=[10, 10, 11]), "lc")
-    >>> resolved = ColumnSpec({"a": Flat("a"), "lc": Nested("lc", "f")}).resolve(nf)
-    >>> batch = resolved.collate(nf)
-    >>> batch["lc"]
-    tensor([[1., 2.],
-            [3., 0.]])
-    >>> batch["lc_mask"]
-    tensor([[ True,  True],
-            [ True, False]])
-    >>> batch["_index"]
-    tensor([10, 11])
-    """
-
-    fields: Mapping[str, Field]
-    index_key: str | None = "_index"
-
-    def resolve(self, schema) -> ResolvedSpec:
-        """Validate the spec against a catalog's schema without computing any data.
-
-        Parameters
-        ----------
-        schema : lsdb.Catalog or pd.DataFrame
-            A catalog (its ``meta`` is used) or any dataframe with the same columns and
-            dtypes, such as ``catalog.meta`` or a computed partition.
-
-        Returns
-        -------
-        ResolvedSpec
-            The validated spec, which can collate frames.
-
-        Raises
-        ------
-        SpecError
-            Listing every problem found.
-        """
-        meta = schema.meta if hasattr(schema, "meta") and not isinstance(schema, pd.DataFrame) else schema
-        if not isinstance(meta, pd.DataFrame):
-            raise TypeError(f"Expected an lsdb catalog or a DataFrame, got {type(schema).__name__}.")
-        problems: list[str] = []
-        if not self.fields:
-            problems.append("The spec has no fields.")
-        resolved: dict[str, ResolvedField] = {}
-        for key, spec_field in self.fields.items():
-            if not isinstance(spec_field, Field):
-                problems.append(f"'{key}': expected a Field such as Flat or Nested, got {type(spec_field)}.")
-                continue
-            result = spec_field.resolve(key, meta, problems)
-            if result is not None:
-                resolved[key] = result
-
-        seen: dict[str, str] = {}
-        if self.index_key is not None:
-            seen[self.index_key] = "the index"
-            if not pd.api.types.is_integer_dtype(meta.index.dtype):
-                problems.append(
-                    f"index_key={self.index_key!r} needs an integer index, but the index has dtype "
-                    f"{meta.index.dtype}; pass index_key=None to disable it."
-                )
-        for key, result in resolved.items():
-            for output_key in result.describe():
-                if output_key in seen:
-                    problems.append(
-                        f"Output key '{output_key}' from '{key}' collides with {seen[output_key]}."
-                    )
-                seen[output_key] = f"'{key}'"
-
-        if problems:
-            hint = ""
-            missing = [p for p in problems if "is not in the catalog" in p]
-            all_columns = getattr(schema, "all_columns", None)
-            if missing and all_columns is not None:
-                hint = (
-                    "\nSome columns may exist on disk but were not loaded; the catalog was opened "
-                    f"with columns {list(meta.columns)}."
-                )
-            raise SpecError("Column spec does not match the catalog:\n  - " + "\n  - ".join(problems) + hint)
-        return ResolvedSpec(resolved, self.index_key)
-
-
-class ResolvedSpec:
-    """A :class:`ColumnSpec` validated against a schema; collates frames into tensors.
-
-    Create it with :meth:`ColumnSpec.resolve`. It is picklable, so it can be shipped to
-    Dask tasks.
-    """
-
-    def __init__(self, fields: dict[str, ResolvedField], index_key: str | None):
-        self.fields = fields
-        self.index_key = index_key
-
-    @property
-    def columns(self) -> list[str]:
-        """Top-level catalog columns needed for collation, in first-use order."""
-        columns: list[str] = []
-        for resolved in self.fields.values():
-            columns.extend(c for c in resolved.columns if c not in columns)
-        return columns
-
-    def describe(self) -> dict[str, TensorInfo]:
-        """Return the dtype and shape of every tensor a batch contains, keyed by output key."""
-        info: dict[str, TensorInfo] = {}
-        for resolved in self.fields.values():
-            info.update(resolved.describe())
-        if self.index_key is not None:
-            info[self.index_key] = TensorInfo(torch.int64, (None,))
-        return info
-
-    def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
-        """Convert the rows of ``frame`` to a batch of tensors.
-
-        Parameters
-        ----------
-        frame : pd.DataFrame
-            Rows to convert, e.g. a slice of a catalog partition. Row order is preserved:
-            row ``i`` of every output tensor comes from row ``i`` of ``frame``.
-
-        Returns
-        -------
-        dict[str, torch.Tensor]
-            One entry per key in :meth:`describe`.
-        """
-        batch: dict[str, torch.Tensor] = {}
-        for resolved in self.fields.values():
-            batch.update(resolved.collate(frame))
-        if self.index_key is not None:
-            batch[self.index_key] = torch.from_numpy(np.array(frame.index.to_numpy(), dtype=np.int64))
-        return batch
diff --git a/tests/lsdb_ml/inference/inference_models.py b/tests/lsdb_ml/inference/inference_models.py
index 3364717..9d2fb0b 100644
--- a/tests/lsdb_ml/inference/inference_models.py
+++ b/tests/lsdb_ml/inference/inference_models.py
@@ -4,8 +4,6 @@
 import torch
 from torch import nn
 
-from lsdb_ml import ColumnSpec, Flat
-
 # Counts factory calls; with processes=False every worker shares this module.
 BUILDS = {"count": 0}
 
@@ -15,14 +13,14 @@ class AffineMag(nn.Module):
 
     def __init__(self):
         super().__init__()
-        self.linear = nn.Linear(1, 1, dtype=torch.float64)
+        self.linear = nn.Linear(1, 1)
         with torch.no_grad():
             self.linear.weight.fill_(2.0)
             self.linear.bias.fill_(1.0)
 
-    def forward(self, mag):
-        """Map (B,) magnitudes to (B,) predictions."""
-        return self.linear(mag.unsqueeze(-1)).squeeze(-1)
+    def forward(self, batch):
+        """Map the (B,) "mag" tensor to (B,) predictions."""
+        return self.linear(batch["mag"].unsqueeze(-1)).squeeze(-1)
 
 
 def counting_factory():
@@ -34,28 +32,29 @@ def counting_factory():
 class MaskedMean(nn.Module):
     """Mean of the real (unpadded) light-curve values; 0 for empty light curves."""
 
-    def forward(self, flux, mask):
+    def forward(self, batch):
         """Average padded values using the mask."""
+        flux, mask = batch["lc.flux"], batch["lc._mask"]
         return (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
 
 
 class Stacked(nn.Module):
     """Two outputs per row, as a (B, 2) tensor."""
 
-    def forward(self, mag):
+    def forward(self, batch):
         """Return [mag, -mag]."""
-        return torch.stack([mag, -mag], dim=1)
+        return torch.stack([batch["mag"], -batch["mag"]], dim=1)
 
 
 class DropsRow(nn.Module):
     """A broken model that returns one prediction too few."""
 
-    def forward(self, mag):
+    def forward(self, batch):
         """Drop the last row."""
-        return mag[:-1]
+        return batch["mag"][:-1]
 
 
-MAG_SPEC = ColumnSpec({"mag": Flat("mag", dtype=torch.float64)})
+MAG_INPUTS = "mag"
 
 
 def check_alignment(result, source_frame):
diff --git a/tests/lsdb_ml/inference/test_predict.py b/tests/lsdb_ml/inference/test_predict.py
index 554616a..1b5efba 100644
--- a/tests/lsdb_ml/inference/test_predict.py
+++ b/tests/lsdb_ml/inference/test_predict.py
@@ -5,7 +5,7 @@
 from distributed import Client, LocalCluster
 from inference_models import (
     BUILDS,
-    MAG_SPEC,
+    MAG_INPUTS,
     AffineMag,
     DropsRow,
     MaskedMean,
@@ -13,8 +13,9 @@
     check_alignment,
     counting_factory,
 )
+from torch import nn
 
-from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested, OutputSpec, predict
+from lsdb_ml import InputError, ModelDeployment, ModelSpec, OutputSpec, predict
 
 
 @pytest.fixture(scope="module")
@@ -41,7 +42,7 @@ def test_index_aligned_after_nontrivial_layout(client, scrambled_catalog, nested
     assert 0 in [len(p.compute(progress_bar=False)) for p in scrambled_catalog.partitions]
 
     with ModelDeployment(ModelSpec(AffineMag)) as deployed:
-        lazy = deployed.predict(scrambled_catalog, MAG_SPEC, {"pred": "float64"}, batch_size=3)
+        lazy = deployed.predict(scrambled_catalog, MAG_INPUTS, {"pred": "float64"}, batch_size=3)
         result = lazy.compute(progress_bar=False)
     assert len(result) == len(computed_input)
     assert result.index.equals(computed_input.index)
@@ -56,7 +57,7 @@ def test_alignment_with_duplicate_index(client, nested_frame):
     frame["ra"] = 10.0  # every object at the same position: one index value repeated
     frame["dec"] = 10.0
     catalog = lsdb.from_dataframe(frame, partition_rows=100)
-    result = predict(catalog, AffineMag(), MAG_SPEC, {"pred": "float64"}, batch_size=7).compute(
+    result = predict(catalog, AffineMag(), MAG_INPUTS, {"pred": "float64"}, batch_size=7).compute(
         progress_bar=False
     )
     assert not result.index.is_unique
@@ -68,10 +69,10 @@ def test_model_loaded_once_per_worker(client, nested_catalog):
     BUILDS["count"] = 0
     spec = ModelSpec(counting_factory)
     with ModelDeployment(spec) as deployed:
-        deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
-        deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+        deployed.predict(nested_catalog, MAG_INPUTS, {"pred": "float64"}).compute(progress_bar=False)
+        deployed.predict(nested_catalog, MAG_INPUTS, {"pred": "float64"}).compute(progress_bar=False)
         # Deploying the same spec again reuses the installed plugin.
-        predict(nested_catalog, spec, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+        predict(nested_catalog, spec, MAG_INPUTS, {"pred": "float64"}).compute(progress_bar=False)
     assert nested_catalog.npartitions > len(client.scheduler_info()["workers"])
     assert BUILDS["count"] == len(client.scheduler_info()["workers"]) == 2
 
@@ -79,7 +80,7 @@ def test_model_loaded_once_per_worker(client, nested_catalog):
 def test_meta_matches_computed(client, nested_catalog):
     """The lazy result's meta has the same columns and dtypes as the computed data."""
     for append in (True, False):
-        lazy = predict(nested_catalog, AffineMag(), MAG_SPEC, {"pred": "float32"}, append=append)
+        lazy = predict(nested_catalog, AffineMag(), MAG_INPUTS, {"pred": "float32"}, append=append)
         computed = lazy.compute(progress_bar=False)
         assert list(lazy.columns) == list(computed.columns)
         assert lazy.dtypes.to_dict() == computed.dtypes.to_dict()
@@ -89,14 +90,13 @@ def test_meta_matches_computed(client, nested_catalog):
 
 
 def test_nested_inputs(client, nested_catalog):
-    """Nested columns are padded and masked, so a masked mean matches a per-row mean."""
-    spec = ColumnSpec({"flux": Nested("lc", "flux", dtype=torch.float64)})
+    """Nested sub-columns arrive padded with a shared mask, so a masked mean matches a per-row mean."""
     result = predict(
         nested_catalog,
         MaskedMean(),
-        spec,
+        "lc.flux",
         {"mean_flux": "float64"},
-        inputs=["flux", "flux_mask"],
+        float_dtype=torch.float64,
         batch_size=9,
     ).compute(progress_bar=False)
     expected = [
@@ -106,9 +106,21 @@ def test_nested_inputs(client, nested_catalog):
     np.testing.assert_allclose(result["mean_flux"].to_numpy(), expected, rtol=1e-6)
 
 
+def test_forward_adapter(client, nested_catalog):
+    """forward= adapts models that take positional tensors instead of the batch dict."""
+    result = predict(
+        nested_catalog,
+        nn.Identity(),
+        ["mag", "lc.flux"],
+        {"pred": "float64"},
+        forward=lambda model, batch: 2 * model(batch["mag"]) + 1,
+    ).compute(progress_bar=False)
+    np.testing.assert_allclose(result["pred"], 2 * result["mag"].astype(np.float64) + 1, rtol=1e-6)
+
+
 def test_multi_column_output(client, nested_catalog):
     """A (B, k) output fills k columns in order."""
-    result = predict(nested_catalog, Stacked(), MAG_SPEC, ["plus", "minus"], append=False).compute(
+    result = predict(nested_catalog, Stacked(), MAG_INPUTS, ["plus", "minus"], append=False).compute(
         progress_bar=False
     )
     np.testing.assert_allclose(result["plus"], -result["minus"])
@@ -116,7 +128,7 @@ def test_multi_column_output(client, nested_catalog):
 
 def test_wrong_row_count_fails(client, nested_catalog):
     """A model that drops rows fails loudly instead of misaligning predictions."""
-    lazy = predict(nested_catalog, DropsRow(), MAG_SPEC, {"pred": "float32"})
+    lazy = predict(nested_catalog, DropsRow(), MAG_INPUTS, {"pred": "float32"})
     with pytest.raises(Exception, match="exactly one prediction per input row"):
         lazy.compute(progress_bar=False)
 
@@ -124,13 +136,13 @@ def test_wrong_row_count_fails(client, nested_catalog):
 def test_validation(client, unloadable_catalog):
     """Bad specs, outputs and arguments fail before any data is read."""
     with pytest.raises(ValueError, match="already exist"):
-        predict(unloadable_catalog, AffineMag(), MAG_SPEC, {"mag": "float32"})
-    with pytest.raises(ValueError, match="pass inputs="):
-        predict(unloadable_catalog, AffineMag(), ColumnSpec({"a": Flat("ra"), "b": Flat("dec")}), ["p"])
-    with pytest.raises(ValueError, match="not in the catalog"):
-        predict(unloadable_catalog, AffineMag(), ColumnSpec({"a": Flat("nope")}), ["p"])
+        predict(unloadable_catalog, AffineMag(), MAG_INPUTS, {"mag": "float32"})
+    with pytest.raises(InputError, match="'nope' is not a column"):
+        predict(unloadable_catalog, AffineMag(), ["mag", "nope"], ["p"])
+    with pytest.raises(TypeError, match="forward must be callable"):
+        predict(unloadable_catalog, AffineMag(), MAG_INPUTS, ["p"], forward="mag")
     with pytest.raises(ValueError, match="batch_size"):
-        predict(unloadable_catalog, AffineMag(), MAG_SPEC, ["p"], batch_size=0)
+        predict(unloadable_catalog, AffineMag(), MAG_INPUTS, ["p"], batch_size=0)
     with pytest.raises(TypeError, match="ModelSpec or torch.nn.Module"):
         ModelDeployment("model")
 
@@ -138,13 +150,13 @@ def test_validation(client, unloadable_catalog):
 def test_closed_deployment(client, nested_catalog):
     """After close, the model is gone from the workers and computing fails clearly."""
     deployed = ModelDeployment(ModelSpec(AffineMag, device="cpu", torch_threads=None))
-    lazy = deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float32"})
+    lazy = deployed.predict(nested_catalog, MAG_INPUTS, {"pred": "float32"})
     deployed.close()
     deployed.close()
     with pytest.raises(Exception, match="is not deployed"):
         lazy.compute(progress_bar=False)
     with pytest.raises(RuntimeError, match="closed"):
-        deployed.predict(nested_catalog, MAG_SPEC, {"pred": "float32"})
+        deployed.predict(nested_catalog, MAG_INPUTS, {"pred": "float32"})
 
 
 def test_weights_loaded_on_workers(client, nested_catalog, tmp_path):
@@ -154,7 +166,7 @@ def test_weights_loaded_on_workers(client, nested_catalog, tmp_path):
         trained.linear.bias.fill_(5.0)
     torch.save(trained.state_dict(), tmp_path / "state.pt")
     spec = ModelSpec(AffineMag, weights=tmp_path / "state.pt")
-    result = predict(nested_catalog, spec, MAG_SPEC, {"pred": "float64"}).compute(progress_bar=False)
+    result = predict(nested_catalog, spec, MAG_INPUTS, {"pred": "float64"}).compute(progress_bar=False)
     np.testing.assert_allclose(result["pred"], 2 * result["mag"].astype(np.float64) + 5, rtol=1e-6)
 
 
diff --git a/tests/lsdb_ml/test_batch.py b/tests/lsdb_ml/test_batch.py
new file mode 100644
index 0000000..25adb92
--- /dev/null
+++ b/tests/lsdb_ml/test_batch.py
@@ -0,0 +1,181 @@
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from lsdb_ml import InputError, plan_inputs
+from lsdb_ml.collate import pad_ragged
+
+
+def reference_rows(frame, column, sub):
+    """Per-row python reference: one array of sub-column values per row."""
+    rows = []
+    for value in frame[column].array:
+        rows.append(value[sub].to_numpy(np.float64) if isinstance(value, pd.DataFrame) else np.zeros(0))
+    return rows
+
+
+def reference_pad(rows, width, keep="head"):
+    """Pad reference rows to ``width`` the slow way."""
+    padded = np.zeros((len(rows), width))
+    mask = np.zeros((len(rows), width), dtype=bool)
+    for i, row in enumerate(rows):
+        row = row[:width] if keep == "head" else row[max(len(row) - width, 0) :]
+        padded[i, : len(row)] = row
+        mask[i, : len(row)] = True
+    return padded, mask
+
+
+@pytest.fixture
+def frame(nested_catalog):
+    """A frame mixing a sliced partition, a whole partition, and a concatenation."""
+    parts = [nested_catalog.partitions[i].compute(progress_bar=False) for i in range(3)]
+    return pd.concat([parts[0].iloc[3:], parts[1], parts[2].iloc[:5]])
+
+
+# Validation against meta
+
+
+def test_unloadable_catalog_fixture_raises(unloadable_catalog):
+    """Guard: the no-compute fixture really fails when a partition is computed."""
+    with pytest.raises(Exception, match="no data should be read"):
+        unloadable_catalog.partitions[0].compute(progress_bar=False)
+
+
+def test_plan_from_catalog_without_compute(unloadable_catalog):
+    """Flat vs nested and dtypes are inferred from meta; nothing is computed."""
+    plan = plan_inputs(unloadable_catalog, ["mag", "id", "flag", "lc.t", "lc.flux"])
+    assert plan.describe() == {
+        "mag": torch.float32,
+        "id": torch.int64,
+        "flag": torch.bool,
+        "lc.t": torch.float32,
+        "lc.flux": torch.float32,
+        "lc._mask": torch.bool,
+        "lc._lengths": torch.int64,
+    }
+    batch = plan.to_tensors(unloadable_catalog.meta)  # the empty meta frame is a valid dry run
+    assert {k: v.shape[0] for k, v in batch.items()} == dict.fromkeys(plan.describe(), 0)
+
+
+def test_all_problems_reported_together(unloadable_catalog):
+    """Every problem is reported in one error."""
+    with pytest.raises(InputError) as err:
+        plan_inputs(unloadable_catalog, ["nope", "lc", "name", "lc.nope", "ra.x", "lc.band"], max_length=0)
+    message = str(err.value)
+    for expected in [
+        "'nope' is not a column",
+        "'lc' is a nested column; name its sub-columns",
+        "'name' has non-numeric dtype",
+        "has no sub-column 'nope'",
+        "'ra' is not a nested column",
+        "'lc.band' has non-numeric type",
+        "max_length must be a positive integer",
+        "others may exist on disk",
+    ]:
+        assert expected in message
+
+
+def test_other_argument_errors(unloadable_catalog):
+    """Bad truncate/float_dtype values, empty inputs and bad schemas are rejected."""
+    with pytest.raises(InputError, match="truncate must be"):
+        plan_inputs(unloadable_catalog, "mag", truncate="middle")
+    with pytest.raises(InputError, match="float_dtype"):
+        plan_inputs(unloadable_catalog, "mag", float_dtype=torch.int32)
+    with pytest.raises(InputError, match="No input columns"):
+        plan_inputs(unloadable_catalog, [])
+    with pytest.raises(TypeError, match="Expected an lsdb catalog"):
+        plan_inputs(42, "mag")
+
+
+def test_reserved_mask_name():
+    """A real sub-column named like the generated mask is rejected."""
+    nf = npd.NestedFrame({"a": [1.0]}, index=[0]).join_nested(
+        npd.NestedFrame({"_mask": [1.0]}, index=[0]), "lc"
+    )
+    with pytest.raises(InputError, match="reserved"):
+        plan_inputs(nf, ["lc._mask"])
+
+
+# Conversion
+
+
+def test_flat_columns(frame):
+    """Flat columns become (B,) tensors with inferred dtypes, in row order."""
+    batch = plan_inputs(frame, ["ra", "mag", "id", "flag"]).to_tensors(frame)
+    np.testing.assert_allclose(batch["ra"].numpy(), frame["ra"].to_numpy(np.float32))
+    assert batch["ra"].dtype == torch.float32 and batch["ra"].shape == (len(frame),)
+    np.testing.assert_array_equal(batch["id"].numpy(), frame["id"].to_numpy())
+    assert batch["flag"].dtype == torch.bool
+    assert all(t.numpy().flags.writeable for t in batch.values())
+
+
+def test_float_dtype(frame):
+    """float_dtype controls floating-point columns only."""
+    batch = plan_inputs(frame, ["ra", "id", "lc.t"], float_dtype=torch.float64).to_tensors(frame)
+    assert batch["ra"].dtype == torch.float64 and batch["lc.t"].dtype == torch.float64
+    assert batch["id"].dtype == torch.int64
+
+
+def test_nested_matches_reference(frame):
+    """Padding matches a per-row reference, including empty and null nested rows."""
+    batch = plan_inputs(frame, ["lc.t", "lc.flux"], float_dtype=torch.float64).to_tensors(frame)
+    lengths = np.array([len(r) for r in reference_rows(frame, "lc", "t")])
+    assert (lengths == 0).any(), "fixture should include empty / null nested rows"
+    for sub in ("t", "flux"):
+        expected, expected_mask = reference_pad(reference_rows(frame, "lc", sub), lengths.max())
+        np.testing.assert_allclose(batch[f"lc.{sub}"].numpy(), expected)
+        np.testing.assert_array_equal(batch["lc._mask"].numpy(), expected_mask)
+    np.testing.assert_array_equal(batch["lc._lengths"].numpy(), lengths)
+
+
+@pytest.mark.parametrize("keep", ["head", "tail"])
+def test_truncation(frame, keep):
+    """Rows over max_length keep their first or last values."""
+    batch = plan_inputs(frame, "lc.t", float_dtype=torch.float64, max_length=2, truncate=keep).to_tensors(
+        frame
+    )
+    rows = reference_rows(frame, "lc", "t")
+    expected, expected_mask = reference_pad(rows, 2, keep=keep)
+    np.testing.assert_allclose(batch["lc.t"].numpy(), expected)
+    np.testing.assert_array_equal(batch["lc._mask"].numpy(), expected_mask)
+    np.testing.assert_array_equal(batch["lc._lengths"].numpy(), np.minimum([len(r) for r in rows], 2))
+
+
+def test_truncation_raises_by_default(frame):
+    """Rows over max_length raise unless a truncate policy is given."""
+    with pytest.raises(ValueError, match="over max_length=2"):
+        plan_inputs(frame, "lc.t", max_length=2).to_tensors(frame)
+
+
+def test_all_rows_empty_nested(frame):
+    """A batch where every nested row is empty pads to width 0."""
+    empty_rows = frame[np.asarray(frame["lc"].nest.len()) == 0]
+    batch = plan_inputs(frame, "lc.flux").to_tensors(empty_rows)
+    assert batch["lc.flux"].shape == (len(empty_rows), 0)
+    assert not batch["lc._mask"].any()
+
+
+def test_missing_values_raise():
+    """Missing flat values and missing values inside nested lists raise."""
+    nf = npd.NestedFrame({"a": [1.0, None, 3.0]}, index=[0, 1, 2])
+    nf = nf.join_nested(npd.NestedFrame({"f": [1.0, None, 2.0]}, index=[0, 0, 2]), "lc")
+    with pytest.raises(ValueError, match="1 row\\(s\\) with missing values"):
+        plan_inputs(nf, "a").to_tensors(nf)
+    with pytest.raises(ValueError, match="missing values"):
+        plan_inputs(nf, "lc.f").to_tensors(nf)
+
+
+def test_row_order_preserved(frame):
+    """Row i of every tensor comes from row i of the frame, even for a shuffled frame."""
+    shuffled = frame.sample(frac=1, random_state=1)
+    batch = plan_inputs(shuffled, ["id", "lc.t"], float_dtype=torch.float64).to_tensors(shuffled)
+    np.testing.assert_array_equal(batch["id"].numpy(), shuffled["id"].to_numpy())
+    np.testing.assert_array_equal(batch["lc._lengths"].numpy(), np.asarray(shuffled["lc"].nest.len()))
+
+
+def test_pad_ragged_length_mismatch():
+    """Inconsistent lengths and values are rejected."""
+    with pytest.raises(ValueError, match="lengths sum to 3"):
+        pad_ragged(np.zeros((2, 1)), np.array([1, 2]))
diff --git a/tests/lsdb_ml/test_collate.py b/tests/lsdb_ml/test_collate.py
deleted file mode 100644
index b0123be..0000000
--- a/tests/lsdb_ml/test_collate.py
+++ /dev/null
@@ -1,136 +0,0 @@
-import nested_pandas as npd
-import numpy as np
-import pandas as pd
-import pytest
-import torch
-
-from lsdb_ml import ColumnSpec, Flat, Nested
-from lsdb_ml.collate import pad_ragged
-
-
-def reference_rows(frame, column, fields):
-    """Per-row python reference: a list of (length, n_fields) arrays."""
-    rows = []
-    for value in frame[column].array:
-        if value is None or not isinstance(value, pd.DataFrame):
-            rows.append(np.zeros((0, len(fields))))
-        else:
-            rows.append(value[fields].to_numpy(dtype=np.float64))
-    return rows
-
-
-def reference_pad(rows, width, keep="head"):
-    """Pad reference rows to ``width`` the slow way."""
-    padded = np.zeros((len(rows), width, rows[0].shape[1] if rows else 0))
-    mask = np.zeros((len(rows), width), dtype=bool)
-    for i, row in enumerate(rows):
-        row = row[:width] if keep == "head" else row[max(len(row) - width, 0) :]
-        padded[i, : len(row)] = row
-        mask[i, : len(row)] = True
-    return padded, mask
-
-
-@pytest.fixture
-def frame(nested_catalog):
-    """A frame mixing a sliced partition, a whole partition, and a concatenation."""
-    # A computed partition, plus a slice and a concatenation, all exercise the
-    # flat-values/offsets bookkeeping on non-trivial arrow layouts.
-    parts = [nested_catalog.partitions[i].compute(progress_bar=False) for i in range(3)]
-    return pd.concat([parts[0].iloc[3:], parts[1], parts[2].iloc[:5]])
-
-
-def test_flat_columns(frame):
-    """Flat columns stack into (B,) or (B, F) tensors with the requested dtype."""
-    resolved = ColumnSpec(
-        {"x": Flat(["ra", "dec", "mag"]), "mag": Flat("mag"), "flag": Flat("flag", dtype=torch.bool)}
-    ).resolve(frame)
-    batch = resolved.collate(frame)
-    assert batch["x"].shape == (len(frame), 3) and batch["x"].dtype == torch.float32
-    np.testing.assert_allclose(batch["x"].numpy(), frame[["ra", "dec", "mag"]].to_numpy(dtype=np.float32))
-    assert batch["mag"].shape == (len(frame),)
-    assert batch["flag"].dtype == torch.bool
-    np.testing.assert_array_equal(batch["flag"].numpy(), frame["flag"].to_numpy(dtype=bool))
-    np.testing.assert_array_equal(batch["_index"].numpy(), frame.index.to_numpy())
-    assert all(t.numpy().flags.writeable for t in batch.values())
-
-
-@pytest.mark.parametrize("fields", [["t", "flux"], ["flux"]])
-def test_nested_matches_reference(frame, fields):
-    """Vectorized nested padding matches a per-row reference, including empty rows."""
-    resolved = ColumnSpec({"lc": Nested("lc", fields, dtype=torch.float64)}).resolve(frame)
-    batch = resolved.collate(frame)
-    rows = reference_rows(frame, "lc", fields)
-    lengths = np.array([len(r) for r in rows])
-    assert (lengths == 0).any(), "fixture should include empty / null nested rows"
-    expected, expected_mask = reference_pad(rows, lengths.max())
-    np.testing.assert_allclose(batch["lc"].numpy(), expected)
-    np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask)
-    np.testing.assert_array_equal(batch["lc_lengths"].numpy(), lengths)
-
-
-def test_nested_single_field_is_2d(frame):
-    """A single nested sub-column produces (B, L) values."""
-    batch = ColumnSpec({"f": Nested("lc", "flux")}).resolve(frame).collate(frame)
-    assert batch["f"].ndim == 2
-    assert batch["f"].shape == batch["f_mask"].shape
-
-
-@pytest.mark.parametrize("keep", ["head", "tail"])
-def test_truncation(frame, keep):
-    """Rows over max_length keep their first or last values."""
-    resolved = ColumnSpec({"lc": Nested("lc", ["t"], max_length=2, truncate=keep, dtype=torch.float64)})
-    batch = resolved.resolve(frame).collate(frame)
-    rows = reference_rows(frame, "lc", ["t"])
-    expected, expected_mask = reference_pad(rows, 2, keep=keep)
-    np.testing.assert_allclose(batch["lc"].numpy(), expected)
-    np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask)
-    np.testing.assert_array_equal(batch["lc_lengths"].numpy(), np.minimum([len(r) for r in rows], 2))
-
-
-def test_truncation_raises_by_default(frame):
-    """Rows over max_length raise unless a truncate policy is given."""
-    with pytest.raises(ValueError, match="over max_length=2"):
-        ColumnSpec({"lc": Nested("lc", "t", max_length=2)}).resolve(frame).collate(frame)
-
-
-def test_empty_frame(frame):
-    """Zero-row frames collate to zero-row tensors."""
-    resolved = ColumnSpec({"x": Flat(["ra", "dec"]), "lc": Nested("lc", ["t", "flux"])}).resolve(frame)
-    batch = resolved.collate(frame.iloc[:0])
-    assert batch["x"].shape == (0, 2)
-    assert batch["lc"].shape == (0, 0, 2)
-    assert batch["lc_mask"].shape == (0, 0)
-    assert batch["_index"].shape == (0,)
-
-
-def test_all_rows_empty_nested(frame):
-    """A batch where every nested row is empty pads to width 0."""
-    lengths = np.asarray(frame["lc"].nest.len())
-    empty_rows = frame[lengths == 0]
-    batch = ColumnSpec({"lc": Nested("lc", "flux")}).resolve(frame).collate(empty_rows)
-    assert batch["lc"].shape == (len(empty_rows), 0)
-    assert not batch["lc_mask"].any()
-
-
-def test_null_values_raise():
-    """Missing flat values and missing values inside nested lists raise."""
-    nf = npd.NestedFrame({"a": [1.0, None, 3.0]}, index=[0, 1, 2])
-    nf = nf.join_nested(npd.NestedFrame({"f": [1.0, None, 2.0]}, index=[0, 0, 2]), "lc")
-    with pytest.raises(ValueError, match="1 row\\(s\\) with missing values"):
-        ColumnSpec({"a": Flat("a")}).resolve(nf).collate(nf)
-    with pytest.raises(ValueError, match="missing values"):
-        ColumnSpec({"lc": Nested("lc", "f")}).resolve(nf).collate(nf)
-
-
-def test_pad_ragged_length_mismatch():
-    """Inconsistent lengths and values are rejected."""
-    with pytest.raises(ValueError, match="lengths sum to 3"):
-        pad_ragged(np.zeros((2, 1)), np.array([1, 2]))
-
-
-def test_row_order_preserved(frame):
-    """Row i of every tensor comes from row i of the frame, even for a shuffled frame."""
-    shuffled = frame.sample(frac=1, random_state=1)
-    batch = ColumnSpec({"id": Flat("id", dtype=torch.int64)}).resolve(shuffled).collate(shuffled)
-    np.testing.assert_array_equal(batch["id"].numpy(), shuffled["id"].to_numpy())
-    np.testing.assert_array_equal(batch["_index"].numpy(), shuffled.index.to_numpy())
diff --git a/tests/lsdb_ml/test_spec.py b/tests/lsdb_ml/test_spec.py
deleted file mode 100644
index 1eb8faa..0000000
--- a/tests/lsdb_ml/test_spec.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import pytest
-import torch
-
-from lsdb_ml import ColumnSpec, Flat, Nested, SpecError, TensorInfo
-
-
-def test_resolve_valid_spec_from_catalog_without_compute(unloadable_catalog):
-    """A valid spec resolves against a catalog without computing any partition."""
-    spec = ColumnSpec(
-        {
-            "x": Flat(["ra", "dec", "mag"]),
-            "flag": Flat("flag", dtype=torch.bool),
-            "lc": Nested("lc", ["t", "flux"], max_length=4, truncate="tail"),
-            "flux": Nested("lc", "flux", max_length=8, truncate="head"),
-        }
-    )
-    resolved = spec.resolve(unloadable_catalog)
-    assert resolved.columns == ["ra", "dec", "mag", "flag", "lc"]
-    assert resolved.describe() == {
-        "x": TensorInfo(torch.float32, (None, 3)),
-        "flag": TensorInfo(torch.bool, (None,)),
-        "lc": TensorInfo(torch.float32, (None, None, 2)),
-        "lc_mask": TensorInfo(torch.bool, (None, None)),
-        "lc_lengths": TensorInfo(torch.int64, (None,)),
-        "flux": TensorInfo(torch.float32, (None, None)),
-        "flux_mask": TensorInfo(torch.bool, (None, None)),
-        "flux_lengths": TensorInfo(torch.int64, (None,)),
-        "_index": TensorInfo(torch.int64, (None,)),
-    }
-
-
-def test_unloadable_catalog_fixture_raises(unloadable_catalog):
-    """Guard: the no-compute fixture really fails when a partition is computed."""
-    with pytest.raises(Exception, match="no data should be read"):
-        unloadable_catalog.partitions[0].compute(progress_bar=False)
-
-
-def test_resolve_accepts_meta_frame(nested_catalog):
-    """resolve() also accepts the meta frame directly."""
-    resolved = ColumnSpec({"x": Flat("ra")}).resolve(nested_catalog.meta)
-    assert resolved.columns == ["ra"]
-
-
-def test_all_problems_reported_together(unloadable_catalog):
-    """Every problem in a spec is reported in one error."""
-    spec = ColumnSpec(
-        {
-            "missing": Flat("nope"),
-            "nested_as_flat": Flat("lc"),
-            "string": Flat("name"),
-            "flat_as_nested": Nested("ra", "x"),
-            "missing_sub": Nested("lc", ["flux", "nope"]),
-            "string_sub": Nested("lc", "band"),
-            "bad_trunc": Nested("lc", "flux", truncate="middle"),
-            "bad_len": Nested("lc", "flux", max_length=0),
-        }
-    )
-    with pytest.raises(SpecError) as err:
-        spec.resolve(unloadable_catalog)
-    message = str(err.value)
-    for expected in [
-        "column 'nope' is not in the catalog",
-        "column 'lc' is nested",
-        "'name' has non-numeric dtype",
-        "column 'ra' is not nested",
-        "has no sub-column 'nope'",
-        "'lc.band' has non-numeric type",
-        "truncate must be",
-        "max_length must be a positive integer",
-    ]:
-        assert expected in message
-    assert "not loaded" in message  # hint for missing columns on a catalog
-
-
-def test_key_collisions(unloadable_catalog):
-    """Output keys may not collide with each other or the index key."""
-    with pytest.raises(SpecError, match="'lc_mask' from 'lc' collides"):
-        ColumnSpec({"lc_mask": Flat("ra"), "lc": Nested("lc", "flux")}).resolve(unloadable_catalog)
-    with pytest.raises(SpecError, match="'_index' from '_index' collides with the index"):
-        ColumnSpec({"_index": Flat("ra")}).resolve(unloadable_catalog)
-    ColumnSpec({"_index": Flat("ra")}, index_key=None).resolve(unloadable_catalog)
-
-
-def test_non_field_and_empty_spec(unloadable_catalog):
-    """Non-Field values and empty specs are rejected."""
-    with pytest.raises(SpecError, match="expected a Field"):
-        ColumnSpec({"x": "ra"}).resolve(unloadable_catalog)
-    with pytest.raises(SpecError, match="no fields"):
-        ColumnSpec({}).resolve(unloadable_catalog)
-
-
-def test_unsupported_dtype(unloadable_catalog):
-    """Torch dtypes without a numpy equivalent are rejected."""
-    with pytest.raises(SpecError, match="no numpy equivalent"):
-        ColumnSpec({"x": Flat("ra", dtype=torch.bfloat16)}).resolve(unloadable_catalog)
-
-
-def test_non_integer_index(nested_frame):
-    """The index key requires an integer index, and can be disabled."""
-    frame = nested_frame.reset_index(drop=True).set_index("name")
-    with pytest.raises(SpecError, match="needs an integer index"):
-        ColumnSpec({"x": Flat("ra")}).resolve(frame)
-    ColumnSpec({"x": Flat("ra")}, index_key=None).resolve(frame)
-
-
-def test_bad_schema_type():
-    """resolve() rejects objects that are neither catalogs nor frames."""
-    with pytest.raises(TypeError, match="Expected an lsdb catalog"):
-        ColumnSpec({"x": Flat("ra")}).resolve(42)

From 33fc2aa6adcf97f991344fd46ded562a9c9d5d9b Mon Sep 17 00:00:00 2001
From: Doug Branton <dougrbranton@gmail.com>
Date: Fri, 16 Oct 2026 14:48:02 -0700
Subject: [PATCH 7/7] Allow max_length and truncate per nested column

Both accept a single value for every nested input or a mapping keyed
by nested column name, for multi-modal catalogs with several light
curve and spectrum columns. Keys are column names rather than dotted
sub-column names because sub-columns share one mask; unknown keys,
dotted keys and invalid values are reported with the other input
problems before any data is read.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/design/mvp-design.md               |  5 +-
 src/lsdb_ml/batch.py                    | 87 +++++++++++++++++++------
 src/lsdb_ml/inference/predict.py        | 18 +++--
 tests/lsdb_ml/inference/test_predict.py | 34 ++++++++++
 tests/lsdb_ml/test_batch.py             | 70 ++++++++++++++++++++
 5 files changed, 187 insertions(+), 27 deletions(-)

diff --git a/docs/design/mvp-design.md b/docs/design/mvp-design.md
index cbaced4..4280582 100644
--- a/docs/design/mvp-design.md
+++ b/docs/design/mvp-design.md
@@ -85,7 +85,10 @@ with ModelDeployment(ModelSpec(MyModel, weights="model.pt")) as deployed:
 - **Batch-level padding is always correct, but has two costs.**
   - *Width follows the longest row.* One very long light curve makes its whole batch that
     wide, so memory and compute scale with `B × L_max` per batch. `max_length` with
-    `truncate` caps this.
+    `truncate` caps this. Both take one value for every nested input, or a mapping per nested
+    column for multi-modal catalogs, e.g. `max_length={"lc_g": 200, "spectrum": 4000}` and
+    `truncate={"lc_g": "tail", "spectrum": "head"}`. Keys are nested column names, not
+    dotted sub-column names, because a column's sub-columns share one mask.
   - *Width varies from batch to batch.* Models must accept a variable `L`. That's natural
     for sequence models and pooling, but a model that needs a fixed `L` (e.g. an MLP over a
     flattened sequence) isn't supported yet. A fixed-width padding option (`pad_to`, or a
diff --git a/src/lsdb_ml/batch.py b/src/lsdb_ml/batch.py
index de0d392..5031b99 100644
--- a/src/lsdb_ml/batch.py
+++ b/src/lsdb_ml/batch.py
@@ -20,7 +20,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 
 import numpy as np
@@ -31,6 +31,8 @@
 
 from lsdb_ml.collate import TruncatePolicy, frame_to_array, pad_ragged
 
+_TRUNCATE_POLICIES = ("raise", "head", "tail")
+
 MASK_SUFFIX = "._mask"
 LENGTHS_SUFFIX = "._lengths"
 
@@ -84,12 +86,13 @@ def __init__(
         flat: list[_Column],
         nested: dict[str, list[_Column]],
         float_dtype: torch.dtype,
-        max_length: int | None,
-        truncate: TruncatePolicy,
+        max_length: dict[str, int | None],
+        truncate: dict[str, TruncatePolicy],
     ):
         self.flat = flat
         self.nested = nested
         self.float_dtype = float_dtype
+        #: Per nested column: longest sequence to keep (None for no limit), and the policy.
         self.max_length = max_length
         self.truncate = truncate
 
@@ -122,8 +125,8 @@ def to_tensors(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
         Raises
         ------
         ValueError
-            If a value is missing, or a nested row is longer than ``max_length`` with
-            ``truncate="raise"``.
+            If a value is missing, or a nested row is longer than its column's
+            ``max_length`` with ``truncate="raise"``.
         """
         batch: dict[str, torch.Tensor] = {}
         for column in self.flat:
@@ -143,7 +146,11 @@ def to_tensors(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]:
                     flat_values[[sub.name]], _numpy_dtype(sub.kind, self.float_dtype), repr(key)
                 )
                 padded, mask, kept = pad_ragged(
-                    values, lengths, max_length=self.max_length, truncate=self.truncate, what=repr(key)
+                    values,
+                    lengths,
+                    max_length=self.max_length[parent],
+                    truncate=self.truncate[parent],
+                    what=repr(key),
                 )
                 batch[key] = torch.from_numpy(np.ascontiguousarray(padded[..., 0]))
             batch[parent + MASK_SUFFIX] = torch.from_numpy(mask)
@@ -155,13 +162,44 @@ def _numpy_dtype(kind: str, float_dtype: torch.dtype) -> np.dtype:
     return torch.empty(0, dtype=_torch_dtype(kind, float_dtype)).numpy().dtype
 
 
+def _valid_max_length(value) -> bool:
+    return value is None or (isinstance(value, int) and not isinstance(value, bool) and value >= 1)
+
+
+def _per_column(option, value, default, nested_columns, valid, describe, problems) -> dict:
+    """Expand a single value or a ``{nested column: value}`` mapping to every nested input."""
+    if not isinstance(value, Mapping):
+        if not valid(value):
+            problems.append(f"{option} must be {describe}, got {value!r}.")
+        return dict.fromkeys(nested_columns, value)
+    resolved = dict.fromkeys(nested_columns, default)
+    for key, item in value.items():
+        parent = str(key).partition(".")[0]
+        if key in nested_columns:
+            if valid(item):
+                resolved[key] = item
+            else:
+                problems.append(f"{option}['{key}'] must be {describe}, got {item!r}.")
+        elif parent in nested_columns:
+            problems.append(
+                f"{option} key '{key}': use the nested column name '{parent}'; its sub-columns share "
+                "one mask, so they are truncated together."
+            )
+        else:
+            problems.append(
+                f"{option} key '{key}' is not a nested column among the inputs "
+                f"(nested inputs: {sorted(nested_columns)})."
+            )
+    return resolved
+
+
 def plan_inputs(
     schema,
     inputs: str | Sequence[str],
     *,
     float_dtype: torch.dtype = torch.float32,
-    max_length: int | None = None,
-    truncate: TruncatePolicy = "raise",
+    max_length: int | Mapping[str, int] | None = None,
+    truncate: TruncatePolicy | Mapping[str, TruncatePolicy] = "raise",
 ) -> InputPlan:
     """Check input column names against a catalog's schema, without computing data.
 
@@ -173,11 +211,15 @@ def plan_inputs(
         Column names: ``"mag"`` for a flat column, ``"lc.flux"`` for a nested sub-column.
     float_dtype : torch.dtype, default torch.float32
         Dtype for floating-point columns.
-    max_length : int, optional
-        Longest nested sequence to keep; see ``truncate``.
-    truncate : {"raise", "head", "tail"}, default "raise"
+    max_length : int or Mapping[str, int], optional
+        Longest nested sequence to keep; see ``truncate``. An int applies to every nested
+        input. A mapping sets it per nested column, keyed by the column name (``"lc"``,
+        not ``"lc.flux"``, since sub-columns share one mask); unlisted columns have no
+        limit.
+    truncate : {"raise", "head", "tail"} or Mapping[str, str], default "raise"
         For nested rows longer than ``max_length``: raise, keep the first values, or keep
-        the last values.
+        the last values. A mapping sets it per nested column, like ``max_length``;
+        unlisted columns use "raise".
 
     Returns
     -------
@@ -210,12 +252,6 @@ def plan_inputs(
     problems: list[str] = []
     if not names:
         problems.append("No input columns were given.")
-    if truncate not in ("raise", "head", "tail"):
-        problems.append(f"truncate must be 'raise', 'head' or 'tail', got {truncate!r}.")
-    if max_length is not None and (
-        isinstance(max_length, bool) or not isinstance(max_length, int) or max_length < 1
-    ):
-        problems.append(f"max_length must be a positive integer, got {max_length!r}.")
     if float_dtype not in (torch.float16, torch.float32, torch.float64):
         problems.append(f"float_dtype must be torch.float16, float32 or float64, got {float_dtype}.")
 
@@ -249,6 +285,19 @@ def plan_inputs(
         else:
             nested.setdefault(parent, []).append(_Column(sub, _kind(meta.dtypes[parent].column_dtypes[sub])))
 
+    lengths_limits = _per_column(
+        "max_length", max_length, None, nested, _valid_max_length, "a positive integer or None", problems
+    )
+    policies = _per_column(
+        "truncate",
+        truncate,
+        "raise",
+        nested,
+        lambda v: v in _TRUNCATE_POLICIES,
+        "'raise', 'head' or 'tail'",
+        problems,
+    )
+
     for parent, subs in nested.items():
         for suffix in (MASK_SUFFIX, LENGTHS_SUFFIX):
             if any(f"{parent}.{sub.name}" == parent + suffix for sub in subs):
@@ -261,4 +310,4 @@ def plan_inputs(
         if any("is not a column" in p for p in problems) and getattr(schema, "all_columns", None) is not None:
             hint = f"\nThe catalog was loaded with columns {list(meta.columns)}; others may exist on disk."
         raise InputError("Input columns do not match the catalog:\n  - " + "\n  - ".join(problems) + hint)
-    return InputPlan(flat, nested, float_dtype, max_length, truncate)
+    return InputPlan(flat, nested, float_dtype, lengths_limits, policies)
diff --git a/src/lsdb_ml/inference/predict.py b/src/lsdb_ml/inference/predict.py
index de00d79..d0300c9 100644
--- a/src/lsdb_ml/inference/predict.py
+++ b/src/lsdb_ml/inference/predict.py
@@ -89,8 +89,8 @@ def predict(
         batch_size: int = 4096,
         append: bool = True,
         float_dtype: torch.dtype = torch.float32,
-        max_length: int | None = None,
-        truncate: TruncatePolicy = "raise",
+        max_length: int | Mapping[str, int] | None = None,
+        truncate: TruncatePolicy | Mapping[str, TruncatePolicy] = "raise",
     ):
         """Return a lazy catalog with the model's predictions for every row.
 
@@ -146,8 +146,8 @@ def predict(
     batch_size: int = 4096,
     append: bool = True,
     float_dtype: torch.dtype = torch.float32,
-    max_length: int | None = None,
-    truncate: TruncatePolicy = "raise",
+    max_length: int | Mapping[str, int] | None = None,
+    truncate: TruncatePolicy | Mapping[str, TruncatePolicy] = "raise",
     client=None,
 ):
     """Run a model over every row of a catalog, returning a lazy catalog of predictions.
@@ -184,12 +184,16 @@ def predict(
         same index).
     float_dtype : torch.dtype, default torch.float32
         Dtype for floating-point input columns.
-    max_length : int, optional
+    max_length : int or Mapping[str, int], optional
         Longest nested sequence to keep; see ``truncate``. Nested inputs are padded to the
         longest row in each batch, so one very long row makes the whole batch that wide.
-    truncate : {"raise", "head", "tail"}, default "raise"
+        An int applies to every nested input; a mapping sets it per nested column, keyed
+        by column name, e.g. ``{"lc_g": 200, "spectrum": 4000}`` (unlisted columns have no
+        limit).
+    truncate : {"raise", "head", "tail"} or Mapping[str, str], default "raise"
         For nested rows longer than ``max_length``: raise, keep the first values, or keep
-        the last values.
+        the last values. A mapping sets it per nested column; unlisted columns use
+        "raise".
     client : distributed.Client, optional
         Client to deploy to. Defaults to the active client; one is required.
 
diff --git a/tests/lsdb_ml/inference/test_predict.py b/tests/lsdb_ml/inference/test_predict.py
index 1b5efba..27a18d0 100644
--- a/tests/lsdb_ml/inference/test_predict.py
+++ b/tests/lsdb_ml/inference/test_predict.py
@@ -118,6 +118,40 @@ def test_forward_adapter(client, nested_catalog):
     np.testing.assert_allclose(result["pred"], 2 * result["mag"].astype(np.float64) + 1, rtol=1e-6)
 
 
+def test_per_column_max_length(client, nested_frame):
+    """predict passes per-column limits through: each nested input is truncated separately."""
+    import lsdb
+
+    rng = np.random.default_rng(2)
+    lengths = rng.integers(10, 30, len(nested_frame))
+    spec = pd.DataFrame(
+        {"flux": rng.normal(0, 1, lengths.sum())}, index=np.repeat(nested_frame.index, lengths)
+    )
+    catalog = lsdb.from_dataframe(nested_frame.join_nested(spec, "spec"), partition_rows=100)
+
+    def widths(model, batch):
+        rows = batch["mag"].shape[0]
+        return torch.stack(
+            [
+                torch.full((rows,), batch["lc.flux"].shape[1]),
+                torch.full((rows,), batch["spec.flux"].shape[1]),
+            ],
+            dim=1,
+        )
+
+    result = predict(
+        catalog,
+        nn.Identity(),
+        ["mag", "lc.flux", "spec.flux"],
+        {"lc_width": "int64", "spec_width": "int64"},
+        forward=widths,
+        max_length={"lc": 2, "spec": 12},
+        truncate="tail",
+        append=False,
+    ).compute(progress_bar=False)
+    assert result["lc_width"].max() <= 2 and result["spec_width"].max() == 12
+
+
 def test_multi_column_output(client, nested_catalog):
     """A (B, k) output fills k columns in order."""
     result = predict(nested_catalog, Stacked(), MAG_INPUTS, ["plus", "minus"], append=False).compute(
diff --git a/tests/lsdb_ml/test_batch.py b/tests/lsdb_ml/test_batch.py
index 25adb92..0c7c908 100644
--- a/tests/lsdb_ml/test_batch.py
+++ b/tests/lsdb_ml/test_batch.py
@@ -179,3 +179,73 @@ def test_pad_ragged_length_mismatch():
     """Inconsistent lengths and values are rejected."""
     with pytest.raises(ValueError, match="lengths sum to 3"):
         pad_ragged(np.zeros((2, 1)), np.array([1, 2]))
+
+
+# Several nested columns
+
+
+@pytest.fixture(scope="module")
+def two_nested(nested_frame):
+    """The test frame plus a second nested column ("spec") with longer rows (10-29 values)."""
+    rng = np.random.default_rng(1)
+    lengths = rng.integers(10, 30, len(nested_frame))
+    spec = pd.DataFrame(
+        {"wave": rng.uniform(4000, 9000, lengths.sum()), "flux": rng.normal(0, 1, lengths.sum())},
+        index=np.repeat(nested_frame.index, lengths),
+    )
+    return nested_frame.join_nested(spec, "spec")
+
+
+def test_per_column_max_length_and_truncate(two_nested):
+    """max_length and truncate can be set per nested column, keyed by column name."""
+    plan = plan_inputs(
+        two_nested,
+        ["lc.t", "lc.flux", "spec.wave"],
+        float_dtype=torch.float64,
+        max_length={"lc": 2, "spec": 12},
+        truncate={"lc": "head", "spec": "tail"},
+    )
+    batch = plan.to_tensors(two_nested)
+    for column, sub, width, keep in [("lc", "t", 2, "head"), ("spec", "wave", 12, "tail")]:
+        expected, expected_mask = reference_pad(reference_rows(two_nested, column, sub), width, keep=keep)
+        np.testing.assert_allclose(batch[f"{column}.{sub}"].numpy(), expected)
+        np.testing.assert_array_equal(batch[f"{column}._mask"].numpy(), expected_mask)
+    # Sub-columns of one nested column are truncated together and share the mask.
+    assert batch["lc.flux"].shape == batch["lc.t"].shape == batch["lc._mask"].shape
+
+
+def test_unlisted_columns_are_unlimited(two_nested):
+    """Columns missing from the max_length mapping keep every value."""
+    batch = plan_inputs(
+        two_nested, ["lc.t", "spec.wave"], max_length={"spec": 12}, truncate="tail"
+    ).to_tensors(two_nested)
+    assert batch["spec.wave"].shape[1] == 12
+    assert batch["lc.t"].shape[1] == int(np.asarray(two_nested["lc"].nest.len()).max())
+
+
+def test_single_value_applies_to_every_nested_column(two_nested):
+    """An int max_length and a single policy apply to all nested inputs."""
+    batch = plan_inputs(two_nested, ["lc.t", "spec.wave"], max_length=3, truncate="head").to_tensors(
+        two_nested
+    )
+    assert batch["lc.t"].shape[1] == batch["spec.wave"].shape[1] == 3
+
+
+def test_per_column_key_errors(two_nested):
+    """Mapping keys must be nested input columns; dotted keys get a hint; values are checked."""
+    with pytest.raises(InputError) as err:
+        plan_inputs(
+            two_nested,
+            ["lc.t", "spec.wave"],
+            max_length={"lc.t": 5, "lcc": 5, "spec": 0},
+            truncate={"lc": "middle", "mag": "head"},
+        )
+    message = str(err.value)
+    for expected in [
+        "max_length key 'lc.t': use the nested column name 'lc'",
+        "max_length key 'lcc' is not a nested column among the inputs",
+        "max_length['spec'] must be a positive integer",
+        "truncate['lc'] must be 'raise', 'head' or 'tail'",
+        "truncate key 'mag' is not a nested column among the inputs",
+    ]:
+        assert expected in message