Skip to content

Registering Custom Operators

HeliaAOT now uses an explicit RegistryContext model.

You build a context once and pass it into AotConverter. The context contains all runtime registries used during conversion:

  • context.litert_parsers
  • context.aot_operator_classes
  • context.transforms
  • context.air_options

Model

  • Registries are context-scoped and explicit.
  • Custom extensions are applied through customizers when building the context.

Default Built-In Flow

CLI builds the default context automatically and passes it to the converter.

For Python API usage, do the same explicitly:

from helia_aot.registry.context import build_default_registry_context
from helia_aot.converter import AotConverter

registry_context = build_default_registry_context()
converter = AotConverter(config, registry_context=registry_context)
converter.convert()

Custom Operator Registration via Customizer

Use a customizer callable that mutates the context before conversion:

from helia_aot.registry.context import RegistryContext, build_default_registry_context
from helia_aot.converter import AotConverter

from my_plugin.ops import parse_custom_fft, CustomFftOperator


def customize_registry(context: RegistryContext) -> None:
    context.litert_parsers.register("CUSTOM_FFT", parse_custom_fft, overwrite=True)
    context.aot_operator_classes.register("CUSTOM_FFT", CustomFftOperator, overwrite=True)


registry_context = build_default_registry_context(
    customizers=[customize_registry],
    allow_override=True,
)

AotConverter(config, registry_context=registry_context).convert()

Use one canonical key string for all custom-op registration points: - LiteRT parser key in context.litert_parsers - AIR op_type emitted by your parser - AOT operator class key in context.aot_operator_classes

For LiteRT custom ops, construct this key from customCode with:

from helia_aot.air import custom_code_to_op_key

Extending Transforms and AIR Options

The same context can be extended beyond parser/operator mappings:

  • context.transforms: add or replace transform classes.
  • context.air_options: add or replace operator options classes.

Key Points

  • Use build_default_registry_context(...) to start from built-ins.
  • Use customizers for custom ops and advanced overrides.
  • Set allow_override=True explicitly when replacing built-ins.
  • Operator/tensor YAML attribute rules (config.operators, memory.tensors) remain separate from registry selection.

Starter Example

For an almost complete, copy/paste starting point (LiteRT parser + AOT operator + customizer), see: