ExTorch.Export.OpRegistry (extorch v0.4.0)

Copy Markdown

Registry for custom ExTorch.Export.OpHandler modules.

Maps op target strings to handler modules so ExTorch.Export can dispatch to external op implementations at both compile time (load) and runtime (forward).

Handlers are registered in two ways:

  1. Config (loaded at first access):

    config :extorch, :export_op_handlers, [MyApp.VisionOps]

  2. Runtime:

    ExTorch.Export.OpRegistry.register(MyApp.VisionOps)

The registry is backed by a persistent_term for fast concurrent reads.

Summary

Functions

lookup(op_target)

Look up the handler module for a given op target string.

register(handler)

Register an ExTorch.Export.OpHandler module.

reset()

Remove all registered handlers and reset the dispatch table.

table()

Returns the full dispatch table as a %{op_target => handler_module} map.

Functions

lookup(op_target)

@spec lookup(String.t()) :: {:ok, module()} | :error

Look up the handler module for a given op target string.

Returns {:ok, handler_module} or :error.

register(handler)

@spec register(module()) :: :ok

Register an ExTorch.Export.OpHandler module.

Merges the handler's ops into the dispatch table. If an op is already registered by another handler, the new handler takes precedence.

reset()

@spec reset() :: :ok

Remove all registered handlers and reset the dispatch table.

table()

@spec table() :: %{required(String.t()) => module()}

Returns the full dispatch table as a %{op_target => handler_module} map.