Optimizer state is derived from the optimizer contract and parameter list:
SGD has none; momentum has one velocity per parameter; AdamW has first and
second moments per parameter plus one global step counter.
242def deriveLearnerStateSchema =
243 (lambda unrestricted contract : (family LearnerContract) .
244 (lambda unrestricted parameters : (family StdList (family ParameterSchema)) .
245 (eliminate
246 LearnerContract
247 (lambda unrestricted current : (family LearnerContract) . (family LearnerStateSchema))
248 contract
249 (branch
250 LearnerSGD
251 .
252 (constructor
253 LearnerStateSchema
254 LearnerStateSchemaOf
255 (constructor StdList StdListEmpty (family LearnerStateSlot))))
256 (branch
257 LearnerMomentum
258 .
259 (constructor
260 LearnerStateSchema
261 LearnerStateSchemaOf
262 (learnerSlotsForKind (constructor LearnerStateKind LearnerVelocity) parameters)))
263 (branch
264 LearnerAdamW
265 .
266 (constructor
267 LearnerStateSchema
268 LearnerStateSchemaOf
269 (stdListAppend
270 (family LearnerStateSlot)
271 (learnerSlotsForKind (constructor LearnerStateKind LearnerFirstMoment) parameters)
272 (stdListAppend
273 (family LearnerStateSlot)
274 (learnerSlotsForKind (constructor LearnerStateKind LearnerSecondMoment) parameters)
275 (constructor
276 StdList
277 StdListCons
278 (family LearnerStateSlot)
279 (constructor
280 LearnerStateSlot
281 LearnerGlobalState
282 (constructor LearnerStateKind LearnerStepCounter))
283 (constructor StdList StdListEmpty (family LearnerStateSlot))))))))))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.