module Coppelius.Learner

import Model.Architecture
import Representation.Schema
import Accelerator.SM86.Instruction

-- Coppelius keeps the objective beside the model geometry, but learning is a
-- separate semantic choice.  The shared Representation.Schema owner defines
-- AdamW; this module only binds that reusable learner to Coppelius.
family CoppeliusLearningSemantics : Type 0
constructor CoppeliusLearningSemanticsValue
field unrestricted coppeliusLearningObjective : (family ModelObjective)
field unrestricted coppeliusLearningLearner : (family LearnerContract)

end-family

def coppeliusObjectiveIdentity : Bytes =
  b"exact-vocabulary-cross-entropy"

def coppeliusLearnerIdentity : Bytes =
  b"adamw"

def coppeliusObjective : (family ModelObjective) =
  (constructor ModelObjective ModelExactVocabulary)

def coppeliusLearner : (family LearnerContract) =
  (constructor LearnerContract LearnerAdamW)

def coppeliusLearningSemantics : (family CoppeliusLearningSemantics) =
  (record
    CoppeliusLearningSemantics
    (coppeliusLearningObjective = coppeliusObjective)
    (coppeliusLearningLearner = coppeliusLearner))

-- AdamW's hyperparameters, exact (numerator, denominator): the plan's
-- per-step scalars (Coppelius.Build.Graph) are computed from them and
-- rounded to binary32 once
def coppeliusLearningRateNumerator : Nat = 1
def coppeliusLearningRateDenominator : Nat = 10000
def coppeliusBeta1Numerator : Nat = 9
def coppeliusBeta1Denominator : Nat = 10
def coppeliusBeta2Numerator : Nat = 999
def coppeliusBeta2Denominator : Nat = 1000
def coppeliusWeightDecayNumerator : Nat = 1
def coppeliusWeightDecayDenominator : Nat = 10
def coppeliusEpsilonNumerator : Nat = 1
def coppeliusEpsilonDenominator : Nat = 100000000
-- The 16-bit format of every half the plan makes -- the matrices' copies,
-- the products' operands, the backward's gradients: bfloat16, binary32's
-- exponent range with an 8-bit significand.  Under binary16 the gradients of
-- the attention scores lay below its subnormals whatever one static scale
-- (a scale that lifted them overflowed layer 0), and the upper layers'
-- queries and keys did not learn; bfloat16 holds them.  The device images
-- are moved to it together (Coppelius.Build.DeviceImages,
-- Accelerator.SM86.HalfFormat).
def coppeliusHalfFormat : (family SM86HalfFormat) =
  (constructor SM86HalfFormat SM86BFloat16)

def coppeliusHalfFormatIdentity : Bytes =
  b"bfloat16"

-- loss scaling: the backward runs on the loss times this, and AdamW's
-- epsilon is scaled by the same factor, so the update is the unscaled
-- loss's: S m / (S sqrt v + S eps).  bfloat16's range needs none (1): the
-- factor binary16 needed (8,192) is what its range lacked.
def coppeliusLossScale : Nat = 1
-- the first update's step number (bias correction counts from 1)
def coppeliusFirstStep : Nat = 1