Concepts

OaModule

The ONE class. nn.Module equivalent. Layers and full models inherit from OaModule. No separate OaModel class.

Two Override Patterns

  • Layers override Forward(OaTensor) — used by OaLinear, OaRMSNorm, etc.
  • Full models override Forward/Backward/Step(OaEngine&) — used by OaLlm

Lifecycle

MethodPurpose
Init()Allocate buffers, load shaders, init weights
Forward()Compute forward pass
Backward()Compute gradients
Step()Optimizer step
Save() / Load().oam checkpoint
Destroy()Free all buffers

Parameter Management

Parameters() returns all trainable tensors (recursive through submodules).RegisterSubmodule() adds child modules for parameter discovery.