mojolearn.SmallMLPTrainer

class mojolearn.SmallMLPTrainer(weight1, bias1, weight2, bias2, *, data_schedule, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01)

Bases: object

Trainer for a fixed 8-16-3 ReLU network with mean cross-entropy and AdamW.

The model is Linear(8, 16) -> ReLU -> Linear(16, 3). Pass weight1 (16, 8), bias1 (16,), weight2 (3, 16) and bias2 (3,) as float32 buffers (NumPy arrays or mojolearn.Array); parameters and inputs are copied. Each batch has 1 to 256 rows of 8 features, with labels 0, 1 or 2. All arithmetic runs on the GPU in the numeric mode selected for the process (identical by default); there is no CPU fallback, and a state saved under one mode loads only under the same mode. Every array returned is a mojolearn.Array.

A step either completes or leaves the state unchanged. No clipping, dropout, mixed precision, automatic differentiation or other layer shapes are supported.

Parameters:

• weight1, bias1, weight2, bias2 (buffers) – Initial parameters, shapes as above.

• data_schedule (dict) – A small JSON-serializable description of your data order (for example a dataset hash, sample order and batch size). It is stored with the step counter and in checkpoints, so a resumed run can continue with the right next batch. It is not validated against the batches you pass; the trainer never fetches or reorders data.

• lr (float, default 0.001) – AdamW learning rate.

• betas (tuple of float, default (0.9, 0.999)) – AdamW moment decay rates.

• eps (float, default 1e-8) – AdamW epsilon.

• weight_decay (float, default 0.01) – AdamW decoupled weight decay.

Examples

>>> import numpy as np, mojolearn
>>> rng = np.random.default_rng(0)
>>> w1 = rng.standard_normal((16, 8)).astype(np.float32) * 0.1
>>> w2 = rng.standard_normal((3, 16)).astype(np.float32) * 0.1
>>> trainer = mojolearn.SmallMLPTrainer(w1, np.zeros(16, np.float32), w2, np.zeros(3, np.float32),
...                                     data_schedule={"batch_size": 32})
>>> X = rng.random((32, 8), dtype=np.float32)
>>> y = rng.integers(0, 3, 32)
>>> out = trainer.train_step(X, y)

__init__(weight1, bias1, weight2, bias2, *, data_schedule, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01)

Methods

apply_gradients(gradients)	Apply one already-combined gradient with AdamW; returns the new step count.
from_checkpoint(path)	Restore a trainer from save_checkpoint output, after checking its schema, size and hash.
load_state_dict(state)	Restore a snapshot from state_dict.
loss_and_grads(X, targets)	Mean cross-entropy loss and parameter gradients of a batch, without a step.
predict_logits(X)	Logits (n_rows, 3) of X (n_rows, 8) with the current weights.
predict_logits_sessions(inputs)	Logits of several independent batches with the current model, concatenated in the order given.
save_checkpoint(path)	Atomically save the full state to path in a bounded, pickle-free MLP checkpoint format.
state_dict()	Return an independent snapshot of the full resumable state (weights, optimizer moments, settings, data schedule).
train_step(X, targets, *[, return_input_grad])	One AdamW training step on a batch.

Attributes

step_	Number of optimizer steps taken.
weights_	weight1, bias1, weight2, bias2.

property step_

Number of optimizer steps taken.

property weights_

weight1, bias1, weight2, bias2.

Type:

Current parameters as a dict

state_dict()

Return an independent snapshot of the full resumable state (weights, optimizer moments, settings, data schedule).

load_state_dict(state)

Restore a snapshot from state_dict. The snapshot is validated and copied before any live state changes.

predict_logits_sessions(inputs)

Logits of several independent batches with the current model, concatenated in the order given.

Parameters:

inputs (list of buffers) – float32 batches, each (rows_i, 8).

predict_logits(X)

Logits (n_rows, 3) of X (n_rows, 8) with the current weights.

loss_and_grads(X, targets)

Mean cross-entropy loss and parameter gradients of a batch, without a step.

Use it with apply_gradients to accumulate microbatches yourself.

Returns:

• loss (float)

• gradients (sequence of Arrays) – Gradients of weight1, bias1, weight2, bias2, in that order.

apply_gradients(gradients)

Apply one already-combined gradient with AdamW; returns the new step count.

gradients must hold exactly the four parameter gradients, in the order weight1, bias1, weight2, bias2, otherwise ValueError. The update either completes or leaves the state unchanged.

train_step(X, targets, *, return_input_grad=False)

One AdamW training step on a batch.

Parameters:

• X (buffer, shape (n_rows, 8)) – float32 batch, 1 to 256 rows.

• targets (buffer, shape (n_rows,)) – Labels 0, 1 or 2.

• return_input_grad (bool, default False) – Also return the gradient with respect to X.

Returns:

dict – step (the new step count), loss and logits from before the update, gradients (the four parameter gradients), and input_grad when requested.

save_checkpoint(path)

Atomically save the full state to path in a bounded, pickle-free MLP checkpoint format.

classmethod from_checkpoint(path)

Restore a trainer from save_checkpoint output, after checking its schema, size and hash.