class

AdamScalarCache

AdamScalarCache

source edit

Per-step scalar broadcast cache for Adam-family GPU updates.

Holds the eight broadcast scalars ( $\beta_1$ , $1-\beta_1$ , $\beta_2$ , $1-\beta_2$ , $\epsilon$ , $\eta$ , $1/(1-\beta_1^t)$ , $1/(1-\beta_2^t)$ ) plus the optional weight-decay scalar and the AdamW decay factor as pre-allocated MLX scalar arrays. All entries are rebuilt once per optimiser step at the first parameter and reused across the remaining ~64 parameters of a typical ResNet-18 / transformer step.

Attributes

validbool

true when the cached arrays match the current step / dtype.

dtDtype

Element type the arrays were built for.

for_stepint64

Step number the arrays were built at.

b1, omb1, b2, omb2mlx array (scalar)

\beta_1

1 - \beta_1

\beta_2

1 - \beta_2

broadcasts.

eps_a, lr_a, wd_amlx array (scalar)

\epsilon

\eta

\lambda

broadcasts. wd_a is present only when weight_decay != 0.

inv_bc1, inv_bc2mlx array (scalar)

1 / (1 - \beta_1^t)

and

1 / (1 - \beta_2^t)

wd_factormlx array (scalar)

(1 - \eta \lambda)

, the AdamW multiplicative decay factor.

Notes

Without this cache, adam_step_gpu rebuilt params.size() * 8 MLX scalar arrays per call (~520 allocations for ResNet-18). Each MLX array construction has constant overhead, and the bias-correction factors don't depend on which parameter is being updated, so the scalars are pure step-level invariants.

The cache invalidates when step_count_ advances, the param dtype changes (mixed-precision case), or when set_lr mutates $\eta$ .

Constructors

ctor