module Model.Architecture

import Model.Config
import Model.Word32

-- Hardware-free model-architecture contract shared by named systems. Concrete
-- configurations live with their systems: Model.AlphaER owns the historical
-- Alpha-ER configuration and Coppelius.Model owns Coppelius. Keeping only the
-- contract here prevents either runnable system from depending on the other.
family ModelPositionEncoding : Type 0
constructor ModelLearnedPositions
constructor ModelRotaryPositions

end-family

family ModelNormKind : Type 0
constructor ModelLayerNormalization
constructor ModelRMSNormalization

end-family

family ModelRouteKind : Type 0
constructor ModelPositionRoute
constructor ModelCyclicRoute
constructor ModelTokenHashRoute

end-family

family ModelOptionalWord32 : Type 0
constructor ModelNoWord32
constructor ModelSomeWord32
field unrestricted modelSomeWord32Value : (family ModelWord32)

end-family

family ModelObjective : Type 0
constructor ModelExactVocabulary
constructor ModelSampledVocabulary
field unrestricted modelSampledNegativeCount : (family ModelWord32)

end-family

family ModelConfig : Type 0
constructor ModelConfigValue
field unrestricted modelVocabSize : (family ModelWord32)
field unrestricted modelBlockSize : (family ModelWord32)
field unrestricted modelLayers : (family ModelWord32)
field unrestricted modelWidth : (family ModelWord32)
field unrestricted modelHeads : (family ModelWord32)
field unrestricted modelFfnStoredWidth : (family ModelWord32)
field unrestricted modelExperts : (family ModelWord32)
field unrestricted modelTopK : (family ModelWord32)
field unrestricted modelRoute : (family ModelRouteKind)
field unrestricted modelStackedExperts : (family ModelBoolean)
field unrestricted modelHeadRank : (family ModelOptionalWord32)
field unrestricted modelAttentionRank : (family ModelOptionalWord32)
field unrestricted modelPositionEncoding : (family ModelPositionEncoding)
field unrestricted modelNormKind : (family ModelNormKind)
field unrestricted modelTieEmbeddings : (family ModelBoolean)
field unrestricted modelObjective : (family ModelObjective)

end-family

-- Field projection for `modelBlockSize`, generated from the declaration: the family
-- has one constructor, so this is the unique total projection.
def modelBlockSize =
  (lambda unrestricted value : (family ModelConfig) .
    (eliminate
      ModelConfig
      (lambda unrestricted current : (family ModelConfig) . (family ModelWord32))
      value
      (branch
        ModelConfigValue
        modelVocabSize
        modelBlockSize
        modelLayers
        modelWidth
        modelHeads
        modelFfnStoredWidth
        modelExperts
        modelTopK
        modelRoute
        modelStackedExperts
        modelHeadRank
        modelAttentionRank
        modelPositionEncoding
        modelNormKind
        modelTieEmbeddings
        modelObjective
        .
        modelBlockSize)))