The nine scalars the AdamW SM86 kernel reads, in its ABI order: gradient
scale, beta1, 1 - beta1, beta2, 1 - beta2, the decay multiplier
1 - lr * weightDecay, the bias-corrected step size
lr * sqrt(1 - beta2^t) / (1 - beta1^t), the epsilon scaled by the same
sqrt(1 - beta2^t), and the constant one. `step` counts from 1.
Everything is binary32, the kernel's own precision. 1 - beta2^t loses
about five significant digits to cancellation at small t (beta2^t is close
to one), so the step size carries a relative error near 1e-5; the kernel
applies the same tolerance to its running moments, and the value is what the
learner declares rather than a number computed elsewhere.
70def bobAdamStepScalars =
71 (lambda unrestricted step : Nat .
72 (let unrestricted rootSecond =
73 (f32-sqrt (f32-subtract bobF32One (bobF32Power bobBeta2 step)))
74 in
75 (bytes-builder-build
76 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobF32One))
77 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobBeta1))
78 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One bobBeta1)))
79 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobBeta2))
80 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One bobBeta2)))
81 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One (f32-multiply bobLearningRate bobWeightDecay))))
82 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE
83 (f32-divide (f32-multiply bobLearningRate rootSecond)
84 (f32-subtract bobF32One (bobF32Power bobBeta1 step)))))
85 (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-multiply bobEpsilon rootSecond)))
86 (bytes-builder-chunk (dataFloat32EncodeLE bobF32One)))))))))))))The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.