Base class for all neural network modules.
Every custom model should subclass this and implement forward.
Submodules assigned as attributes are tracked automatically.
Notes
Attribute routing: Setting an attribute follows this priority order:
- If the value is a
lucid.nn.Parameter→ stored in_parameters. - If the value is a
Module→ stored in_modules. - If the name is already a registered buffer and the value is a
lucid.TensororNone→ stays in_buffers. - Otherwise → plain Python attribute.
To register a non-parameter tensor (e.g. a running mean), call
register_buffer explicitly. Rule 3 is what makes the
registration survive: assigning to a buffer updates it rather than
demoting it to a plain attribute, so it stays in state_dict and
keeps moving with to. Assigning something that is neither a
tensor nor None does un-register it, which is the escape hatch.
Examples
>>> import lucid
>>> import lucid.nn as nn
>>> class MLP(nn.Module):
... def __init__(self):
... super().__init__()
... self.fc1 = nn.Linear(10, 20)
... self.fc2 = nn.Linear(20, 1)
...
... def forward(self, x):
... return self.fc2(lucid.relu(self.fc1(x)))
...
>>> model = MLP()
>>> model(lucid.randn(4, 10)).shape
(4, 1)Used by 39
- lucid.compile
- lucid.compile._core.bn_runstats
- lucid.compile._core.fallback
- lucid.compile._core.signature
- lucid.compile._entry.function
- lucid.compile._entry.fused_step
- lucid.compile._entry.module
- lucid.compile._entry.segmented_step
- lucid.compile._entry.step
- lucid.coreml
- lucid.coreml._build
- lucid.coreml._model
… 27 more
Constructors
2Initialise the instance. See the class docstring for parameter semantics.
Run forward with the registered hooks around it.
Always call the module rather than forward directly:
forward pre-hooks, forward hooks and backward hooks only run on
this path. The arguments and the return type are those of the
subclass's own forward.
Instance methods
44Add a child module.
Apply fn recursively to every submodule (including self).
Cast all parameters and buffers to bfloat16.
Yield all buffer tensors.
Yield direct child modules.
No-op compatibility stub.
External codepaths often call model.compile() to opt into JIT
acceleration; Lucid has no such layer, so this returns self
unchanged rather than crashing the caller. Any positional or
keyword arguments are accepted and ignored.
Move all parameters and buffers to CPU.
Cast all parameters and buffers to float64.
Set this module and all children to evaluation mode.
Override to add extra repr info (e.g. Linear shows in_features, etc.).
Cast all parameters and buffers to float32.
Override in subclasses to define the computation.
Return buffer at dotted path, e.g. 'bn.running_mean'.
Return extra state to include in state_dict. Override in subclasses.
Return parameter at dotted path, e.g. 'fc.weight'.
Return submodule at dotted path, e.g. 'encoder.layer.0'.
Cast all parameters and buffers to float16.
Whether any module here is still waiting to see an input.
A lazy layer has no parameters until its first forward, and the
parameters it eventually builds are new objects. Anything that
took a snapshot of parameters before then — an optimiser,
an EMA, a parameter-server shard — holds a list those objects are
not in, and will never touch them again.
load_state_dict
→objectload_state_dict(state_dict: dict[str, Tensor], strict: bool = True, assign: bool = False)Load parameters from a state_dict.
Calls each module's _load_from_state_dict recursively.
Returns _IncompatibleKeys(missing_keys, unexpected_keys) on success.
Raises RuntimeError if strict=True and any keys are missing
or unexpected, or if any error_msgs accumulated during loading.
Parameters
state_dictdictstrictbool= TrueTrue (default) require an exact key match; raise on any
missing or unexpected keys.assignbool= FalseTrue replace each parameter/buffer object with the
loaded tensor, taking its dtype and device; each module's
_load_from_state_dict sees it as
local_metadata["assign_to_params_buffers"]. If False
(default) copy data into the existing parameter preserving its
dtype and device. Either way a shape that differs from the
module's is an error.Move all parameters and buffers to Apple Metal GPU.
Yield this module and all submodules (depth-first).
named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True)Yield (name, buffer) pairs.
Yield (name, child_module) pairs.
named_modules(memo: set[int] | None = None, prefix: str = '', remove_duplicate: bool = True)Yield (name, module) pairs.
named_parameters(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True)Yield (qualified_name, Parameter) pairs from this module's tree.
Parameters
prefixstr= ''"".recursebool= TrueTrue (default), descend into submodules. When
False, yield only this module's directly-attached
parameters.remove_duplicatebool= TrueTrue (default), each unique Parameter
object is yielded only once even if referenced by multiple
attributes — matches the reference framework's contract.Yield all Parameters in this module (and children if recurse=True).
Warns when a lazy layer has not seen an input yet.
The list is complete — a lazy layer registers
lucid.nn.parameter.UninitializedParameter placeholders
at construction, and those objects are the ones the real weights
later occupy, so an optimizer built from this list does train
them. What is not yet true is their shape: they are (0,)
until the first forward, so anything reading shapes, counting
elements or flattening them into a vector is reading nothing.
A warning rather than a refusal because zero_grad and
requires_grad_ legitimately walk a tree in this state. The
test suite promotes it to an error (see filterwarnings in
pyproject.toml), so the gate is strict where a running
program is merely told.
register_backward_hook(hook: _BackwardHook)Deprecated alias for register_full_backward_hook.
register_buffer
→Noneregister_buffer(name: str, tensor: Tensor | None, persistent: bool = True)Register a buffer tensor. Non-persistent buffers are excluded from state_dict.
register_forward_hook(hook: _ForwardHook, prepend: bool = False, with_kwargs: bool = False, always_call: bool = False)Register a hook called after forward().
register_forward_pre_hook(hook: _ForwardPreHook, prepend: bool = False, with_kwargs: bool = False)Register a hook called before forward().
register_full_backward_hook(hook: _BackwardHook, prepend: bool = False)Register a backward hook. Returns a RemovableHandle.
register_full_backward_pre_hook(hook: _BackwardHook, prepend: bool = False)Register a hook to be called before backward hooks.
register_load_state_dict_post_hook(hook: Callable[..., object])Register a post-hook called after this module loads state_dict.
Hook signature: hook(module, incompatible_keys) -> None.
register_load_state_dict_pre_hook(hook: Callable[..., object])Register a pre-hook called when this module loads state_dict.
Hook signature: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None
The hook may mutate state_dict/missing/unexpected/error_msgs.
Alias for add_module.
Register a Parameter under the given name.
Restore extra state loaded from state_dict. Override in subclasses.
No-op on Apple Silicon (unified memory is always shared).
state_dict
→dictstate_dict(destination: dict[str, Tensor] | None = None, prefix: str = '', keep_vars: bool = False)Return an ordered dict mapping parameter / buffer names to tensors.
Includes every learnable parameter, every persistent buffer
(persistent=True at register time), and every nested
submodule's contribution under a dotted-path prefix.
The returned OrderedDict carries a _metadata
attribute mapping module_path → {"version": int} for every
module that defines a _version class attribute.
lucid.save preserves this attribute across disk
round-trips so version-aware _load_from_state_dict hooks
can migrate older checkpoints.
Parameters
destinationdict= NoneNone.prefixstr= ''"".keep_varsbool= FalseTrue keep tensors attached to autograd (return
them as-is); when False (default) detach them so the
state dict is safe to serialise / cross threads.Returns
dictOrderedDict mapping qualified parameter / buffer paths to their tensor values.
Move/cast all parameters and buffers, preserving Parameter object identity.
Takes what lucid.Tensor.to takes: a device, a dtype, both, or
a tensor whose device and dtype to match. A device moves every
tensor. A dtype must be floating point or complex, and it casts the
floating and complex tensors only — an integer or bool buffer such as
BatchNorm.num_batches_tracked keeps its dtype, as in the reference
framework, so a checkpoint round-trip does not change its type.
Raises
TypeError.to forms.Move parameters and buffers to device without copying data.
The reference framework uses to_empty to materialise a
model originally constructed on the meta device. Lucid has no
meta device, so this method exists for API parity and
delegates to the standard to when a device is supplied.
Parameters
deviceobject= Nonedevice, or engine enum).
When None (default) the call is a no-op returning
self.recursebool= Trueto call. Default
True — every child module is moved as well.Returns
SelfThe same module, parameters / buffers now on device.
Set this module and all children to training mode.
Cast every parameter and buffer with lucid.Tensor.type.
Unlike to, this casts the integer and bool buffers too —
BatchNorm.num_batches_tracked included — as the reference does.
Parameters
dst_typelucid.dtype or str"lucid.DoubleTensor" to the CPU,
"lucid.metal.HalfTensor" to Metal, as the reference's names do.Returns
Moduleself, its Parameter objects kept.
Raises
ValueErrordst_type names no legacy type ("float64", float).TypeErrordst_type is neither a dtype, a string nor a type.RuntimeErrordst_type is an integer or bool dtype and a parameter
requires grad, which only a floating or complex tensor can.
Nothing is cast then.Examples
>>> import lucid
>>> import lucid.nn as nn
>>> nn.Linear(2, 2).type("lucid.HalfTensor").weight.dtype
lucid.float16Zero gradients of all parameters.