Tensor lane

The tensor lane provides zero-copy borrow / commit over the Polars frame for kernels that are easier (or faster) to express as array math. Borrow a numeric column as a NumPy view, compute, and commit derived columns back atomically — routed through the snapshot-view contract so the writes stay observable.

Install a packed TensorLane in setup (not inside step): packing rewrites the schema and is visible to contract="check" as a mid-step mutation if done during the step body.

x, is_view = model.agents.borrow("x")   # zero-copy where the layout allows
model.agents.commit(x=x * 2.0)          # atomic write-back, contract-observed

Polars-backed numeric tensor lane.

A thin, version-invariant bridge that exposes a contiguous (N, F) float view over a chosen set of homogeneous numeric agent columns, so that fused elementwise updates and (especially) dense interaction kernels can run as NumPy tensor / einsum operations and write their results back into the model’s Polars DataFrame – which remains the single system of record.

Design goals

• One backing store. The hot numeric attributes live inside Polars (either as one Array(F) column or as F plain columns); the tensor view borrows that memory rather than maintaining a parallel buffer. A borrow is a point-in-time, read-only snapshot – committing a column rebinds its buffer, so outstanding borrows of that column go stale (borrow inputs, then commit).

• Faster only where it helps. Borrowing is most valuable for read-heavy kernels (interactions / reductions) where Polars is forced into an O(N^2) self-join; for trivial elementwise maps, plain Polars columns are already fine and this lane offers no benefit (see the project benchmarks).

Version invariance

This module never branches on polars.__version__ and never touches Polars/Arrow private internals. Instead it probes behaviour once at import time:

• does pl.Series(name, <2-D ndarray>) yield something whose to_numpy() is a 2-D, zero-copy view? -> ARRAY_BACKING_AVAILABLE

• is a single numeric column’s to_numpy() zero-copy? -> SINGLE_COL_ZERO_COPY

Correctness does not depend on either probe succeeding: when the fast path is unavailable the lane falls back to a copy-based path that is correct on every Polars version. Zero-copy is an optimisation that is measured at runtime (via numpy.shares_memory), never assumed. The only public-API calls used – Series.to_numpy() (no kwargs), pl.Series(name, ndarray), DataFrame.select/with_columns/drop/get_column, Series.n_chunks/rechunk – have been stable across the Polars 0.20 -> 1.x line.

class ambr.tensor_lane.TensorLane(model, columns: ~typing.List[str], dtype=<class 'numpy.float64'>, prefer_array: bool = True)

Bases: object

A packed (N, F) numeric view over a model’s homogeneous columns.

Use borrow() to read the packed tensor (zero-copy where the build supports it), commit() to write a full (N, F) result back, and write_result() to attach a derived 1-D column (e.g. an interaction output) without disturbing the pack.

Install the lane in setup (or before run), not mid-step: packing folds named columns into an Array column and will show up as a schema mutation under contract="check". Call unpack() after the run if you need the named columns back for reporting.

The backing strategy is chosen by the capability probe, not by version:

• "array" – the F columns are folded into one Array(F) column; borrow() returns a zero-copy (N, F) view.

• "stacked" – the F columns are kept as-is; borrow() returns a (N, F) copy via DataFrame.to_numpy (correct on every version).

PACKED_SUFFIX = '__packed__'

borrow() → Tuple[ndarray, bool]

Return (arr, is_view) for the packed (N, F) tensor (read-only).

commit(arr: ndarray) → None

Write a full (N, F) result back into the model’s DataFrame.

property is_zero_copy: bool

Whether borrow() returns a view (vs a copy) on this build.

property mode: str

unpack() → None

Restore the F named columns (for relational / named-attribute access).

write_result(name: str, values) → None

Attach a derived 1-D column (e.g. an interaction output).

ambr.tensor_lane.borrow_numeric(model, column: str) → Tuple[ndarray, bool]

Borrow a single numeric column as a 1-D array (read-only).

Returns (array, is_view) where is_view reports whether the array shares memory with the Polars buffer (zero-copy) on this build. The array is always marked read-only, on every build, so the rule that a borrow must not be mutated in place holds regardless of whether the zero-copy fast path was taken (otherwise a stray x += ... would raise on a zero-copy build but silently corrupt a throwaway copy on others).

The array is a point-in-time snapshot: a later commit_columns() (or any write) on this column rebinds the Polars buffer and leaves this array reading the old one. Borrow all inputs first, compute, then commit.

ambr.tensor_lane.commit_columns(model, **columns) → None

Write several 1-D result columns back in a single, atomic with_columns.

Updating existing columns (rather than adding new ones) keeps the agent schema stable across a step – which is what Model.run(contract='check') verifies. Each committed column’s existing dtype is preserved, so a column is never silently re-typed. Pre-declare these columns in setup so the contract certificate stays clean.

Writes go through model._set_frame so the snapshot-view contract can witness the commit and flag same-step duplicate commits of a column.