Python API

Machine learning

The Python-visible autograd, loss, module, optimizer, neural-network, and training surface.

Generated from Source/Python15 functions30 classes1 enums

This inventory is generated from symbols registered by OA's nanobind source and public Python facade. It deliberately excludes unbound C++ internals. Signatures show the Python keyword names declared by each binding; follow the source link for ownership, capability, and execution details.

oa.ml

Functions

Bce(a, b)

BceBwd(a, b)

CrossEntropy(logits, targets)

Cross-entropy loss for classification

CrossEntropyBwd(logits, targets)

Cross-entropy backward: gradient w.r.t. logits

GradEnabled()

IsEnabled()

L1(a, b)

L1Bwd(a, b)

MakeMuonAdamWOptimizer(model, config)

Build official Muon+AdamW composite (Muon on 2D body, AdamW on embed/head/1D)

Mse(a, b)

MseBwd(a, b)

SetEnabled(...)

SetGradEnabled(...)

SmoothL1(a, b)

SmoothL1Bwd(a, b)

Classes

GradientTape4 members

__init__(...)

constructor

Backward(root)

Reverse-mode autodiff from scalar loss root

method

__enter__(...)

method

__exit__(...)

method

OaAdam3 members

__init__(params, lr, beta1, beta2, eps)

method

Step(...)

method

ZeroGrad(...)

method

OaAdamW3 members

__init__(params, lr, beta1, beta2, eps, weight_decay)

AdamW optimizer (decoupled weight decay)

method

Step(...)

Apply one optimizer step (update weights)

method

ZeroGrad(...)

Zero all parameter gradients

method

OaByteEmbedding2 members

__init__(...)

constructor

DModel(...)

prop_ro

OaCallback0 members

OaCbProgressBar2 members

__init__(...)

constructor

AddMetric(metric)

method

OaCbSummary1 members

__init__(...)

constructor

OaCbTraining5 members

OnTrainBegin(...)

method

OnEpochBegin(...)

method

OnStepEnd(...)

method

OnEpochEnd(...)

method

OnTrainEnd(...)

method

OaEmbedding5 members

__init__(...)

constructor

Forward(input)

Forward pass: lookup embeddings for each token index

method

Parameters(...)

Get trainable parameters [embed weight]

method

NumEmbeddings(...)

prop_ro

EmbeddingDim(...)

prop_ro

OaEmpyrealmCore1 members

__init__(...)

constructor

OaGpuTimingStats6 members

Count(...)

ro

MeanMs(...)

ro

MinMs(...)

ro

MedianMs(...)

ro

P95Ms(...)

ro

LastMs(...)

ro

OaGru4 members

__init__(...)

constructor

InputSize(...)

prop_ro

HiddenSize(...)

prop_ro

NumLayers(...)

prop_ro

OaItTraining27 members

__init__(optimizer, config)

method

IsDone(...)

method

Next(...)

method

Next(loss)

method

Reset(...)

method

Index(...)

method

RecordLoss(loss)

method

RecordAccuracy(accuracy)

method

Finish(...)

method

AddCallback(callback)

method

AddMetric(metric)

method

StepCount(...)

method

TotalSteps(...)

method

Epoch(...)

method

StepInEpoch(...)

method

TotalEpochs(...)

method

LiveLoss(...)

method

LastLoss(...)

method

LiveAccuracy(...)

method

ElapsedSeconds(...)

method

LastGpuMs(...)

method

GpuTimingStats(...)

Exact GPU per-step timing (mean/min/median/p95/last ms)

method

WallMsPerStep(...)

method

WallSamplesPerSecond(...)

method

GpuSamplesPerSecond(...)

method

WallUnitsPerSecond(...)

method

GpuUnitsPerSecond(...)

method

OaItTrainingConfig6 members

__init__(...)

constructor

TotalSteps(...)

rw

StepsPerEpoch(...)

rw

BatchSize(...)

rw

SequenceLength(...)

rw

SequenceUnit(...)

rw

OaLayerNorm1 members

__init__(...)

constructor

OaLinear4 members

__init__(...)

constructor

Forward(input)

Forward pass: Out = Input @ W.T + b

method

SetActivation(activation)

Set activation function (None, Relu, Gelu)

method

Parameters(...)

Get trainable parameters [weight, bias]

method

OaMamba3Module1 members

__init__(...)

constructor

OaMetric3 members

Name(...)

method

Result(...)

method

Reset(...)

method

OaMetricAccuracy1 members

__init__(...)

constructor

OaMetricLoss1 members

__init__(name)

method

OaModule8 members

Forward(input)

method

Parameters(...)

Get direct parameters only. Use AllParameterPtrs for nested modules.

method

AllParameterPtrs(...)

Get recursive trainable parameter pointers.

method

NumParameters(...)

method

SetName(...)

method

GetName(...)

method

Save(path)

Serialize module parameters to an .oam file

method

Load(path)

Load module parameters from an .oam file

method

OaMuon3 members

__init__(params, lr, beta, weight_decay, eps, ns5_iters)

Muon optimizer (2D hidden matrices)

method

Step(...)

method

ZeroGrad(...)

method

OaMuonAdamWConfig11 members

__init__(...)

constructor

MuonLr(...)

rw

AdamWLr(...)

rw

MuonBeta(...)

rw

MuonWeightDecay(...)

rw

MuonEps(...)

rw

MuonNs5Iters(...)

rw

AdamWBeta1(...)

rw

AdamWBeta2(...)

rw

AdamWEps(...)

rw

AdamWWeightDecay(...)

rw

OaOptimizer6 members

Step(...)

method

ZeroGrad(...)

method

SetLr(lr)

method

Lr(...)

method

GetLr(...)

method

GetStep(...)

method

OaOptimizerComposite3 members

__init__(...)

constructor

Step(...)

method

ZeroGrad(...)

method

OaOptimizerNoOp1 members

__init__(...)

constructor

OaParameter4 members

Name(...)

prop_ro

Data(...)

prop_rw

Grad(...)

prop_rw

RequiresGrad(...)

rw

OaRnn4 members

__init__(...)

constructor

InputSize(...)

prop_ro

HiddenSize(...)

prop_ro

NumLayers(...)

prop_ro

OaSGD3 members

__init__(params, lr, momentum, weight_decay)

method

Step(...)

method

ZeroGrad(...)

method

OaTransformerBlock2 members

__init__(...)

constructor

SetSeqLen(seq_len)

Update the runtime sequence length without replacing model weights

method

Enums

OaActivationSource

None · Relu · Gelu