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

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
class Request(rf.RequestBase):
    type: Literal["signed_scalar"] = "signed_scalar"
    clip: Annotated[float, pydantic.Field(gt=0.0)] = 10.0

Register components with their exact names. 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.

Custom atoms supply a physical-type matcher:

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. 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 must accept annotated schema and address parameters. 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.

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) requires annotated parameters. Use batch.trainable to select objectives and target state to determine where content is meaningful. module.track((address, strata, metric, component), value=...) records metrics; 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 runtime adds state, inferred, and optional normalized embedding; extensions must not declare these reserved members. Omit both output hooks when only embeddings are public. See Prediction output.

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. Register Lightning callback classes or factories through extension.callback(...); every rank must use the same registrations.

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.