Custom tensorfields

Create an extension when a field needs a new representation, encoder, objective, or output contract. Use a preprocessor for renaming, filtering, joins, and derived values that existing types can already represent.

The extension owns datatype semantics. Shared data code owns Arrow containers, queries, branch geometry, and masking. Adding a datatype should not require a special case in those shared modules.

Register the request

Suppose a signed balance needs a representation that the existing Number field does not provide. Start with an extension and its declared schema options. This excerpt defines only the request; the runtime components below are also needed.

from typing import Annotated, Literal

import pydantic
import relflow as rf

signed_scalar = rf.Extension(name="signed_scalar", types=(int | float,))

@signed_scalar.register(rf.Component.Request)
class Request(rf.RequestBase):
    type: Literal["signed_scalar"] = "signed_scalar"
    clip: Annotated[float, pydantic.Field(gt=0.0)] = 10.0

Pass the component role explicitly; classes and functions may use any Python name. A separately defined writer registers with signed_scalar.register(rf.Component.write, write_signed_scalar), where write_signed_scalar implements the writing contract below.

Extension names such as signed_scalar must use lowercase letters, digits, and underscores. A duplicate extension name replaces the registry entry with a warning.

Request options inherit shared leaf behavior. Parents supply field names; constructors accept configuration only. Name Pydantic validators check_...; use post_bind_validate() when validation needs the bound tree.

Declare Arrow compatibility

types describes accepted terminal value families, using Python type names for their Arrow equivalents. Each tuple entry is a distinct family:

types=(int | float,) permits numeric promotion within a field.
types=(str, bytes) accepts either family but rejects an Arrow union spanning both.
types=(dict,) accepts struct or map atoms; the extension validates their meaning.

List containers, dictionary encodings, and unions are traversed to their terminal types. Null is a shared state and needs no type entry. These checks see Arrow types after ingress conversion, not the original Python objects. Semantic shape checks, such as vector width, belong in the codec.

A custom atom, such as an exact-decimal monetary value, supplies a physical-type matcher. The extension still owns its numerical representation:

from decimal import Decimal

import pyarrow as pa

decimal = rf.Extension(
    name="decimal",
    types=(Decimal,),
    arrow={Decimal: pa.types.is_decimal},
)

Each matcher accepts a pa.DataType and returns a Boolean. Its key must appear in types. Extension.prepare(...) validates the complete leaf, decodes Arrow dictionary wrappers, preserves directly accepted logical extension types, and safely promotes compatible unions where possible. A remaining union must be handled by the extension codec.

Implement the runtime contract

Component	Contract
`Request`	Pydantic field options and semantic validation.
`TensorField`	Subclass `TensorFieldBase`; convert separate input and target projections to tensors.
`Embedder`	Subclass `EmbedderBase`; turn compact `TensorInput` into a `Parcel` with model-width vectors.
`Decoder`	Subclass `DecoderBase`; turn pooled context into decoded tensors. Required for reconstruction or `embed=True`.
`loss`	Return a scalar tensor for selected targets. Required when a reconstructing mask reaches the leaf.
`output`, `write`	Optional pair declaring and producing extension-owned Arrow output.
`observe`, `learn`	Optional pair for pristine exposure summaries and model-side state updates.

An input-only field needs Request, TensorField, and Embedder. A reconstructing ancestor mask also requires Decoder and loss for every affected active leaf. The schema checks these capabilities when binding, loading, and mutating.

Tensorize and embed

TensorField.new must accept these parameters:

import torch

@classmethod
def new(
    cls,
    input: rf.RaggedField,
    target: rf.RaggedField,
    present: torch.Tensor,
    trainable: torch.Tensor,
    inferred: torch.Tensor,
    address: rf.Address,
    schema: rf.Schema,
    strata: rf.Strata,
    context: rf.Context,
):
    ...

This is a signature excerpt; implementations supply the body. Retain state, content, present, trainable, inferred, and targets on the tensorfield. State follows schema geometry; content may add trailing axes. Fill targets only when the address is in schema.objectives; built-in tensorfields leave it empty otherwise, because nothing reads reconstruction targets for a leaf without an objective. Use the shared rf.Tokens, with separate presence and target-selection bits rather than another masking path.

A RaggedField stores Arrow values, flat int8 state, flat int64 placement, and dense shape. field.dense gives shaped NumPy state; field.place(...) scatters one encoded row per retained value into that geometry. Validate and encode columnarly. Input and target projections already account for queries, overflow, masking, and source-less prediction.

Embedder.forward(inputs) receives only compact present state/content, with no targets or loss bits. Return a Parcel with payload shape (present_count, d_model), Boolean presence shape (present_count,), and the correct leaf-to-parent route. EmbedderBase.embed restores fixed geometry. Parcel and Prediction are available from relflow.structs.packages.

Both Embedder and Decoder constructors receive schema and address. DecoderBase owns query/mean pooling over ancestor context and optional aligned sibling conditioning. Implement decode(pooled) to return a TensorDict of raw predictions. Pass conditioned=False to the base constructor to disable direct sibling conditioning.

After constructing both modules, relflow calls decoder.attach(embedder.context). Override DecoderBase.attach(context) when decoding needs the same live resource as encoding, such as a learned inverse transform. Keep persistent ownership in the embedder and retain a reference to its context in the decoder; checkpoint restoration then updates one authoritative resource.

Observe and learn

An embedder’s context property exposes its worker-safe resource; it becomes context.state in tensorization. context.salt is shared across fields in one encoded batch.

observe(field, *, address, schema, state, learn) sees the canonical pristine field before projection and returns a fixed-schema TensorDict summary or None. Workers may update explicit interprocess resources, not authoritative model modules. learn(module, observation, *, address, strata) applies the carried summary on the model side. The extension owns distributed reduction. With learn=False, observation must leave learned state unchanged.

Loss and output

loss(module, prediction, batch, strata) returns the scalar objective. Use batch.trainable to select objectives and target state to determine where content is meaningful. module.track((address, strata, metric, component), value=...) records metrics using the evaluation naming convention; relflow applies the request’s loss weight.

output(module, address) returns a stable pa.StructType or None. relflow compiles it once, then passes that exact type to write(module, prediction, datatype). The writer returns one flat StructArray coordinate per decoded position, with exactly the declared types and order. It must not infer types from values, filter rows, or wrap repeated axes.

The prediction runtime adds state and inferred to rich leaf results and writes normalized embeddings in a separate representations column. Only state and inferred are reserved within a leaf result; an extension may own an embedding payload member without colliding with the generic representation column. Omit both output hooks when only generic embeddings are public. Every List, LargeList, or FixedSizeList child in a declared output must be named element; declare it explicitly with pa.field("element", child_type, ...). The planner validates this recursively before writing so compliant Parquet preserves nested field names and the declared type.

Packaging and validation

Schema restoration resolves request types against the live registry. Import the extension before constructing or loading its models; registration need not precede importing relflow. Checkpoints store schema and model state, not the extension’s Python implementation. Make the same implementation and its resources available in loading, data-loader, DDP, and serving processes.

Declare optional dependencies with requires={"import_name": "distribution[extra]"}. relflow checks availability before graph construction, including load and rebuild. Declare Lightning callback classes or factories with rf.Extension(..., callbacks=(...)); every rank must declare the same callbacks.

Start from a nearby built-in and its tests: Boolean for a scalar, Number for normalization, Set for vocabularies, or DateParts for structured content. Verify schema round trips, nested axes, null/padded/masked/skipped positions, source-less prediction, compact embedding shapes, finite losses, stable Arrow output, checkpoint state, and the devices/process modes your extension claims to support. Keep the implementation importable and test it with the relflow version you ship.