nexuml.core.base_layer¶
Base pipeline layer definitions for NexuML.
Bases: Enum
Lightning trainer phase reported to pipeline layers.
Bases: Module
Base class for all NexuML pipeline layers.
Subclasses implement forward_tensor and optionally override lifecycle
hooks (on_fit_start, on_fit_end). The base forward method
routes each keys_in key through forward_tensor and writes results
to the corresponding keys_out key in the output TensorDict.
merge_td(x: TensorDict | Tensor) -> torch.Tensor
Concatenate multiple keys_in tensors along the token (dim=1).
Returns:
| Type | Description |
|---|---|
Tensor
|
Concatenated tensor from the TensorDict keys, or x unchanged |
Tensor
|
when it is already a plain tensor. |