ModuleRegistry¶

A generic, type-safe plugin system used throughout Molfun for registering and retrieving modular components (attention mechanisms, blocks, embedders, structure modules, losses).

Quick Start¶

from molfun.modules.registry import ModuleRegistry

# Create a registry for a new component type
MY_REGISTRY = ModuleRegistry("my_component")

# Register a class
@MY_REGISTRY.register("custom")
class MyCustomComponent:
    def __init__(self, dim: int = 64):
        self.dim = dim

# Retrieve and instantiate
cls = MY_REGISTRY.get("custom")
instance = MY_REGISTRY.build("custom", dim=128)

# List all registered
print(MY_REGISTRY.list())  # ["custom"]

Class Reference¶

ModuleRegistry ¶

Name → class registry with optional build-time validation.

Each registry instance is scoped to a single module family (attention, block, structure_module, ...) and can enforce that registered classes inherit from a given base.

init ¶

__init__(name: str, base_class: type | None = None)

Parameters:

Name	Type	Description	Default
`name`	`str`	Human-readable family name (for error messages).	required
`base_class`	`type \| None`	If set, `register()` will reject classes that don't subclass this type.	`None`

register ¶

register(name: str)

Class decorator::

@REGISTRY.register("my_module")
class MyModule(BaseModule): ...

build ¶

build(name: str, **kwargs: Any) -> nn.Module

Instantiate a registered module by name with keyword arguments.

get ¶

get(name: str) -> type[nn.Module] | None

Return the class or None (no KeyError).

list ¶

list() -> list[str]

Return sorted list of registered names.

register¶

Decorator to register a class under a given name.

@MY_REGISTRY.register("custom")
class MyComponent:
    ...

Parameter	Type	Description
`name`	`str`	Unique identifier for this component

Returns: Decorator that registers the class and returns it unchanged.

Raises: ValueError if the name is already registered.

build¶

Instantiate a registered class with the given arguments.

instance = MY_REGISTRY.build("custom", dim=128, dropout=0.1)

Parameter	Type	Description
`name`	`str`	Registered component name
`**kwargs`	`dict`	Arguments passed to the class constructor

Returns: Instance of the registered class.

Raises: KeyError if the name is not registered.

get¶

Retrieve the class (not an instance) by name.

cls = MY_REGISTRY.get("custom")
instance = cls(dim=128)

Parameter	Type	Description
`name`	`str`	Registered component name

Returns: The registered class.

list¶

Return all registered names as a sorted list of strings.

names = MY_REGISTRY.list()
# ["custom", "default", "flash"]

Returns: list[str]

Dunder Methods¶

# Dictionary-style access
cls = MY_REGISTRY["custom"]

# Membership test
if "custom" in MY_REGISTRY:
    ...

# Iteration
for name in MY_REGISTRY:
    print(name)

# Length
print(len(MY_REGISTRY))  # number of registered components

Built-in Registries¶

Molfun ships with several pre-populated registries:

Registry	Module	Contents
`ATTENTION_REGISTRY`	`molfun.modules.attention`	flash, standard, linear, gated
`BLOCK_REGISTRY`	`molfun.modules.blocks`	pairformer, evoformer, simple_transformer
`STRUCTURE_MODULE_REGISTRY`	`molfun.modules.structure_module`	ipa, diffusion
`EMBEDDER_REGISTRY`	`molfun.modules.embedders`	input, esm
`LOSS_REGISTRY`	`molfun.losses`	mse, mae, huber, pearson, openfold