content: {-# LANGUAGE AllowAmbiguousTypes, CPP, ConstraintKinds, DataKinds, DefaultSignatures, DeriveAnyClass, DeriveGeneric, EmptyCase      #-}
  {-# LANGUAGE FlexibleContexts, FlexibleInstances, FunctionalDependencies, GADTs, KindSignatures, MultiParamTypeClasses, MultiWayIf  #-}
  {-# LANGUAGE OverloadedLabels, OverloadedStrings, PartialTypeSignatures, PolyKinds, QuasiQuotes, RankNTypes, ScopedTypeVariables    #-}
  {-# LANGUAGE StandaloneDeriving, TemplateHaskell, TupleSections, TypeApplications, TypeFamilies, TypeFamilyDependencies, TypeInType #-}
  {-# LANGUAGE TypeOperators, UndecidableInstances                                                                                    #-}
  {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
  
  -- | Store tensors and models (collections of tensors) to files.
  --
  -- Files are in PyTorch TorchScript format, a zip file containing metadata and
  -- tensors. When such files are produced by PyTorch they contain the code to run
  -- the model, we neither read this code nor write it to the file when we produce
  -- saved models. A model for us is a collection of weigths, which can be
  -- interpreted some Haskell function.
  
  module Torch.StoredModel where
  import           Control.Monad
  import           Data.IORef
  import qualified Data.Map                 as M
  import           Data.Maybe
  import           Data.ProtoLens.Labels    ()
  import           Data.Singletons
  import           Data.Singletons.TH
  import           Data.Singletons.TypeLits
  import           Data.Text                (Text)
  import qualified Data.Text                as T
  import qualified Data.Vector              as VB
  import           Foreign.ForeignPtr
  import           Generics.Eot             as GE
  import qualified GHC.TypeLits             as TL
  import qualified Torch.C.Types            as C
  import qualified Torch.C.Variable         as C
  import           Torch.Misc
  import           Torch.Tensor             as Tensor
  import           Torch.Types              as Tensor
  
  -- | A stored model, these can be saved and loaded. They should be treated as
  -- opaque.
  data StoredModel = StoredModel (IORef StoredModelData)
  
  data StoredModelData =
    StoredModelData { smodelRoot     :: ForeignPtr C.CJitScriptModule
                    , smodelEntities :: M.Map Text C.ModuleEntityType
                    , smodelTensors  :: M.Map Text (ForeignPtr C.CTensor)
                    }
    deriving (Show)
  
  -- * What can be saved or loaded
  
  -- | Anything that can be saved or loaded in a model is an instance of
  -- 'Stored'. You should not make such instances unless you create new low-level
  -- components, instead, simply derive 'Generic' and 'Stored'.
  --
  -- It's important to know how Stored works. It relies on the TorchScript format,
  -- and making it interoperable with PyTorch. On-disk it's a dictionary that maps
  -- strings to Tensors. PyTorch also stores code in the resulting file, we don't
  -- have a way to do this right now, but have the basic building blocks for
  -- extracting models for production (TODO make this happen one day).
  --
  -- That means every tensor needs a name. 'Stored' provides a few options foree
  -- naming the contents of the dictionary when you add tensors to it. You can
  -- either rely on autogenerated names, that's what `load` and `save` do,
  -- or provide your own names, that's what `loadWithNames` and saveWithNames` do.
  -- The named versions of these functions take as input a series of nested
  -- tuples that matches the shape of the data being passed in, look at some
  -- examples to se how this shakes out.
  --
  -- If you don't provide a name, names will be derived from the structure that
  -- stores the tensors, just as in PyTorch. Except our structures are different
  -- and you will end up with different names. For PyTorch interop you have to provide
  -- names. Probably not a bad idea to do so in general for future-proofing models.
  --
  -- An intermediate version is relying on the Haskell naming but providing
  -- a custom prefix. The naming scheme in Haskell & PyTorch divides
  -- different levels of the hierarchy by a dot. You can save multiple datastructures
  -- that would otherwise clash in the same file by providing each its own unique prefix
  -- so you will end up with names like "mystruct.weights" and "otherstruct.weights".
  --
  -- NB: PyTorch has three model file formats. The original one, ONNX, and
  -- TorchScript. Most models you will find are in the original format, but this
  -- is extremely difficult to deal with. Models can easily be converted to the
  -- newer TorchScript format, the only complication being that you need to look
  -- at the model and provide input tensors of the correct shape. PyTorch will
  -- tell you this if you print the model. I'm sure someone can further automate
  -- this if needed. For example, coverting AlexNet to a TorchScript model:
  -- @
  --    import torch
  --    import torchvision.models as models
  --    model = models.alexnet(pretrained=True)
  --    script_module = torch.jit.trace(model, torch.randn(10, 3, 224, 224))
  --    script_module.save('alexnet.pt')
  --    exit()
  -- @
  -- The on-disk format for TorchScript is a simple and well-organized zip file
  -- if you're curious. We have other functions for dealing with such models, in
  -- particular `listStoredModelTensors`.
  --
  -- We generally prefer the extension .ht for Haskell-Torch, compared to Python's
  -- .pt for PyTorch. Despite the two formats being compatible. We want to
  -- distinguish the two because Haskell-Torch files don't contain code.
  class Stored a where
    -- | Load a stored model from disk. Names will be assigend according to the
    -- Haskell datastructure.
    load :: StoredModel -> IO a
    load sm = loadWithPrefix sm ""
    -- | Load a stored model from disk. Names will be assigend according to the
    -- Haskell datastructure but you get to provide a prefix. Very useful when
    -- storing multiple datastructures in the same file.
    loadWithPrefix :: StoredModel -> Text -> IO a
    default loadWithPrefix :: (HasEot a, EotStored GE.Datatype (Eot a)) => StoredModel -> Text -> IO a
    loadWithPrefix sm prefix = fromEot <$> eotLoadWithPrefix (datatype (Proxy @a)) sm prefix
    -- | Load a stored model from disk with the supplied names. You must provide a
    -- nested tuple in the same shape as your data (look at
    -- @:t namedParameters (Proxy :: Proxy YourTy)@ in ghci).
    -- This allows you to load models from PyTorch and future-proof Haskell models.
    loadWithNames :: ParameterNames a => StoredModel -> NamedParameters a -> IO a
    default loadWithNames :: (ParameterNames a
                            ,HasEot a
                            ,EotStoredWithNames (Eot a)) => StoredModel -> NamedParameters a -> IO a
    loadWithNames sm names = loadWithNamesUnsafe sm (namedParameters (Proxy :: Proxy a) names)
    -- | This is unsafe and internal, you should never use it.
    loadWithNamesUnsafe :: ParameterNames a => StoredModel -> ParameterTree -> IO a
    default loadWithNamesUnsafe :: (HasEot a, EotStoredWithNames (Eot a)) => StoredModel -> ParameterTree -> IO a
    loadWithNamesUnsafe sm names = fromEot <$> eotLoadWithNames sm names
    ------------------------------------------------------------------------
    -- | Save a stored model from disk. Names will be assigend according to the
    -- Haskell datastructure.
    save :: StoredModel -> a -> IO StoredModel
    save sm a = saveWithPrefix sm "" a
    -- | Save a stored model from disk. Names will be assigend according to the
    -- Haskell datastructure but you get to provide a prefix. Very useful when
    -- storing multiple datastructures in the same file.
    saveWithPrefix :: StoredModel -> Text -> a -> IO StoredModel
    default saveWithPrefix :: (HasEot a, EotStored GE.Datatype (Eot a)) => StoredModel -> Text -> a -> IO StoredModel
    saveWithPrefix sm prefix a = eotSaveWithPrefix (datatype (Proxy @a)) sm prefix (toEot a)
    -- | Save a stored model from disk with the supplied names. You must provide a
    -- nested tuple in the same shape as your data (look at
    -- @:t namedParameters (Proxy :: Proxy YourTy)@ in ghci).
    -- This allows you to save models from PyTorch and future-proof Haskell models.
    saveWithNames :: ParameterNames a => StoredModel -> NamedParameters a -> a -> IO StoredModel
    default saveWithNames :: (ParameterNames a
                            ,HasEot a
                            ,EotStoredWithNames (Eot a)) => StoredModel -> NamedParameters a -> a -> IO StoredModel
    saveWithNames sm names a = saveWithNamesUnsafe sm (namedParameters (Proxy :: Proxy a) names) a
    -- | This is unsafe and internal, you should never use it.
    saveWithNamesUnsafe :: ParameterNames a => StoredModel -> ParameterTree -> a -> IO StoredModel
    default saveWithNamesUnsafe :: (HasEot a, EotStoredWithNames (Eot a)) => StoredModel -> ParameterTree -> a -> IO StoredModel
    saveWithNamesUnsafe sm names a = eotSaveWithNames sm names (toEot a)
  
  -- | This allows you to name all of the pieces of a structure saved to a file.
  -- Given a data type it derives a nested structure of tuples, the same shape as
  -- the constructors, bottoming out in predefined network layers. You fill in
  -- those tuples with names which can be used to store and load data from
  -- disk. This may find other uses one day. Unless you're making new layers you
  -- should not make instances of this class. Just derive from `Generic` and
  -- `ParameterNames`.
  class ParameterNames x where
    type NamedParameters x :: *
    type NamedParameters x = EotNamedParameters (Eot x)
    namedParameters :: Proxy x -> NamedParameters x -> ParameterTree
    default namedParameters :: (EotNamedParameters (Eot x) ~ NamedParameters x,
                               HasEot x, EotParameterNames (Eot x)) => Proxy x -> NamedParameters x -> ParameterTree
    namedParameters _ = eotNamedParameters (Proxy @(Eot x))
  
  -- * Internal classes that enable deriving Stored and ParameterNames.
  
  deriving instance (ParameterNames a1, ParameterNames a2)
                  => ParameterNames (a1,a2)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3)
                  => ParameterNames (a1,a2,a3)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4)
                  => ParameterNames (a1,a2,a3,a4)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4, ParameterNames a5)
                  => ParameterNames (a1,a2,a3,a4,a5)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4, ParameterNames a5, ParameterNames a6)
                  => ParameterNames (a1,a2,a3,a4,a5,a6)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4, ParameterNames a5, ParameterNames a6, ParameterNames a7)
                  => ParameterNames (a1,a2,a3,a4,a5,a6,a7)
  -- TODO I would like more instances here, models get big! But we get a
  -- generics-related error:
  --     • Couldn't match type ‘Rep (a1, a2, a3, a4, a5, a6, a7, a8)’
  --                     with ‘M1 D c0 f0’
  --        arising from a use of ‘Torch.StoredModel.$dmnamedParameters’
  --      The type variables ‘c0’, ‘f0’ are ambiguous
  -- What's missng to make this work?
  
  deriving instance (ParameterNames a1, ParameterNames a2
                    ,Stored a1, Stored a2)
                  => Stored (a1,a2)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3
                    ,Stored a1, Stored a2, Stored a3)
                  => Stored (a1,a2,a3)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4
                    ,Stored a1, Stored a2, Stored a3, Stored a4)
                  => Stored (a1,a2,a3,a4)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4, ParameterNames a5
                    ,Stored a1, Stored a2, Stored a3, Stored a4, Stored a5)
                  => Stored (a1,a2,a3,a4,a5)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4, ParameterNames a5, ParameterNames a6
                    ,Stored a1, Stored a2, Stored a3, Stored a4, Stored a5, Stored a6)
                  => Stored (a1,a2,a3,a4,a5,a6)
  deriving instance (ParameterNames a1, ParameterNames a2, ParameterNames a3, ParameterNames a4, ParameterNames a5, ParameterNames a6, ParameterNames a7
                    ,Stored a1, Stored a2, Stored a3, Stored a4, Stored a5, Stored a6, Stored a7)
                  => Stored (a1,a2,a3,a4,a5,a6,a7)
  
  deriving instance ParameterNames (RNNState ty ki batch isBidirectional nrLayers hiddenF)
  deriving instance ParameterNames (RNNStateBatchFirst ty ki batch isBidirectional nrLayers hiddenF)
  deriving instance ParameterNames (LSTMState ty ki batch isBidirectional nrLayers hiddenF)
  deriving instance ParameterNames (LSTMStateBatchFirst ty ki batch isBidirectional nrLayers hiddenF)
  
  deriving instance (TensorConstraints ty ki '[batch], KnownNat (NrOfRNNDirections isBidirectional)
                    ,KnownNat nrLayers, KnownNat batch, KnownNat hiddenF
                    ,ParameterNames (RNNState ty ki batch isBidirectional nrLayers hiddenF))
                  => Stored (RNNState ty ki batch isBidirectional nrLayers hiddenF)
  deriving instance (TensorConstraints ty ki '[batch], KnownNat (NrOfRNNDirections isBidirectional)
                    ,KnownNat nrLayers, KnownNat batch, KnownNat hiddenF
                    ,ParameterNames (RNNState ty ki batch isBidirectional nrLayers hiddenF))
                  => Stored (RNNStateBatchFirst ty ki batch isBidirectional nrLayers hiddenF)
  deriving instance (TensorConstraints ty ki '[batch], KnownNat (NrOfRNNDirections isBidirectional)
                    ,KnownNat nrLayers, KnownNat batch, KnownNat hiddenF
                    ,ParameterNames (RNNState ty ki batch isBidirectional nrLayers hiddenF))
                  => Stored (LSTMState ty ki batch isBidirectional nrLayers hiddenF)
  deriving instance (TensorConstraints ty ki '[batch], KnownNat (NrOfRNNDirections isBidirectional)
                    ,KnownNat nrLayers, KnownNat batch, KnownNat hiddenF
                    ,ParameterNames (RNNState ty ki batch isBidirectional nrLayers hiddenF))
                  => Stored (LSTMStateBatchFirst ty ki batch isBidirectional nrLayers hiddenF)
  
  data ParameterTree = ParameterLeaf [Text] | ParameterBranch [ParameterTree]
    deriving (Show)
  
  class EotParameterNames x where
    type EotNamedParameters x :: *
    eotNamedParameters :: Proxy x -> EotNamedParameters x -> ParameterTree
  
  instance EotParameterNames fields => EotParameterNames (Either fields Void) where
    type EotNamedParameters (Either fields Void) = EotNamedParameters fields
    eotNamedParameters _ = eotNamedParameters (Proxy @fields)
  
  instance ParameterNames x => EotParameterNames (x, ()) where
    type EotNamedParameters (x, ()) = NamedParameters x
    eotNamedParameters _ x' = namedParameters (Proxy @x) x'
  
  instance (ParameterNames x, ParameterNames x1) => EotParameterNames (x, (x1, ())) where
    type EotNamedParameters (x, (x1, ())) = (NamedParameters x, NamedParameters x1)
    eotNamedParameters _ (x',x1') = ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2) => EotParameterNames (x, (x1, (x2, ()))) where
    type EotNamedParameters (x, (x1, (x2, ()))) = (NamedParameters x, NamedParameters x1, NamedParameters x2)
    eotNamedParameters _ (x',x1',x2') = ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3) => EotParameterNames (x, (x1, (x2, (x3, ())))) where
    type EotNamedParameters (x, (x1, (x2, (x3, ())))) = (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3)
    eotNamedParameters _ (x',x1',x2',x3') = ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2', namedParameters (Proxy @x3) x3']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4) => EotParameterNames (x, (x1, (x2, (x3, (x4, ()))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, ()))))) = (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4)
    eotNamedParameters _ (x',x1',x2',x3',x4') = ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2', namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, ())))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, ())))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5, ParameterNames x6)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, ()))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, ()))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5, NamedParameters x6)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, ())))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, ())))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7, ParameterNames x8)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, ()))))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, ()))))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7, NamedParameters x8)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7',x8') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7', namedParameters (Proxy @x8) x8']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7, ParameterNames x8, ParameterNames x9)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, ())))))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, ())))))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7, NamedParameters x8, NamedParameters x9)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7',x8',x9') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7', namedParameters (Proxy @x8) x8'
                      ,namedParameters (Proxy @x9) x9']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7, ParameterNames x8, ParameterNames x9, ParameterNames x10)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, ()))))))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, ()))))))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7, NamedParameters x8, NamedParameters x9, NamedParameters x10)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7',x8',x9',x10') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7', namedParameters (Proxy @x8) x8'
                      ,namedParameters (Proxy @x9) x9', namedParameters (Proxy @x10) x10']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7, ParameterNames x8, ParameterNames x9, ParameterNames x10, ParameterNames x11)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, (x11, ())))))))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, (x11, ())))))))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7, NamedParameters x8, NamedParameters x9, NamedParameters x10, NamedParameters x11)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7',x8',x9',x10',x11') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7', namedParameters (Proxy @x8) x8'
                      ,namedParameters (Proxy @x9) x9', namedParameters (Proxy @x10) x10', namedParameters (Proxy @x11) x11']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7, ParameterNames x8, ParameterNames x9, ParameterNames x10, ParameterNames x11
           ,ParameterNames x12)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, (x11, (x12, ()))))))))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, (x11, (x12, ()))))))))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7, NamedParameters x8, NamedParameters x9, NamedParameters x10, NamedParameters x11
         ,NamedParameters x12)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7',x8',x9',x10',x11',x12') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7', namedParameters (Proxy @x8) x8'
                      ,namedParameters (Proxy @x9) x9', namedParameters (Proxy @x10) x10', namedParameters (Proxy @x11) x11'
                      ,namedParameters (Proxy @x12) x12']
  
  instance (ParameterNames x, ParameterNames x1, ParameterNames x2, ParameterNames x3, ParameterNames x4, ParameterNames x5
           ,ParameterNames x6, ParameterNames x7, ParameterNames x8, ParameterNames x9, ParameterNames x10, ParameterNames x11
           ,ParameterNames x12, ParameterNames x13)
          => EotParameterNames (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, (x11, (x12, (x13, ())))))))))))))) where
    type EotNamedParameters (x, (x1, (x2, (x3, (x4, (x5, (x6, (x7, (x8, (x9, (x10, (x11, (x12, (x13, ())))))))))))))) =
         (NamedParameters x, NamedParameters x1, NamedParameters x2, NamedParameters x3, NamedParameters x4, NamedParameters x5
         ,NamedParameters x6, NamedParameters x7, NamedParameters x8, NamedParameters x9, NamedParameters x10, NamedParameters x11
         ,NamedParameters x12, NamedParameters x13)
    eotNamedParameters _ (x',x1',x2',x3',x4',x5',x6',x7',x8',x9',x10',x11',x12',x13') =
      ParameterBranch [namedParameters (Proxy @x) x', namedParameters (Proxy @x1) x1', namedParameters (Proxy @x2) x2'
                      ,namedParameters (Proxy @x3) x3', namedParameters (Proxy @x4) x4', namedParameters (Proxy @x5) x5'
                      ,namedParameters (Proxy @x6) x6', namedParameters (Proxy @x7) x7', namedParameters (Proxy @x8) x8'
                      ,namedParameters (Proxy @x9) x9', namedParameters (Proxy @x10) x10', namedParameters (Proxy @x11) x11'
                      ,namedParameters (Proxy @x12) x12', namedParameters (Proxy @x13) x13']
  
  instance EotParameterNames () where
    type EotNamedParameters () = ()
    eotNamedParameters _ _ = ParameterLeaf []
  
  instance ParameterNames Int where
    type NamedParameters Int = Text
    namedParameters _ t = ParameterLeaf [t]
  
  instance ParameterNames String where
    type NamedParameters String = (Text,Text)
    namedParameters _ (t,t') = ParameterLeaf [t,t']
  
  instance ParameterNames (Tensor ty ki sz) where
    type NamedParameters (Tensor ty ki sz) = Text
    namedParameters _ t = ParameterLeaf [t]
  
  instance ParameterNames (LinearParam ty ki inF outF) where
    type NamedParameters (LinearParam ty ki inF outF) = (Text, Text)
    namedParameters _ (t,t') = ParameterLeaf [t,t']
  
  instance ParameterNames (ConvParam ty ki outChans otherSz) where
    type NamedParameters (ConvParam ty ki outChans otherSz) = (Text, Text)
    namedParameters _ (t,t') = ParameterLeaf [t,t']
  
  instance ParameterNames (AffineParam ty ki sz) where
    type NamedParameters (AffineParam ty ki sz) = (Text, Text)
    namedParameters _ (t,t') = ParameterLeaf [t,t']
  
  instance ParameterNames (EmbeddingParam ty ki nrEmbeddings embeddingDim) where
    type NamedParameters (EmbeddingParam ty ki nrEmbeddings embeddingDim) = Text
    namedParameters _ t = ParameterLeaf [t]
  
  instance (SingI nrLayers, SingI (NrOfRNNDirections isBidirectional))
         => ParameterNames (GenericRNNParam ty ki gateSize inF hiddenF nrLayers isBidirectional batchFirst) where
    -- TODO Right now the nr of names is dynamic, through the lists below. We
    -- should make this static.
    type NamedParameters (GenericRNNParam ty ki gateSize inF hiddenF nrLayers isBidirectional batchFirst) = ([Text], [Text], [Text], [Text])
    namedParameters _ (t,t',t'',t''') =
      let nrOfTensors = demoteN @nrLayers * demoteN @(NrOfRNNDirections isBidirectional)
      in if | length t    /= nrOfTensors ->
                error $ "Too few RNN names for the weights, wanted " <> show nrOfTensors <> " got " <> show (length t)
            | length t'   /= nrOfTensors ->
                error $ "Too few RNN names for the biases, wanted " <> show nrOfTensors <> " got " <> show  (length t')
            | length t''  /= nrOfTensors ->
                error $ "Too few RNN names for the hidden weights, wanted " <> show nrOfTensors <> " got " <> show  (length t'')
            | length t''' /= nrOfTensors ->
                error $ "Too few RNN names for the hidden biases, wanted " <> show nrOfTensors <> " got " <> show  (length t''')
            | otherwise -> ParameterLeaf (t ++ t' ++ t'' ++ t''')
  
  instance (ParameterNames (GenericRNNParam ty ki 1 inF hiddenF nrLayers isBidirectional batchFirst))
         => ParameterNames (RNNParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) where
    type NamedParameters (RNNParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) =
      NamedParameters (GenericRNNParam ty ki 1 inF hiddenF nrLayers isBidirectional batchFirst)
    namedParameters _ = namedParameters (Proxy :: Proxy (GenericRNNParam ty ki 1 inF hiddenF nrLayers isBidirectional batchFirst))
  
  instance (ParameterNames (GenericRNNParam ty ki 3 inF hiddenF nrLayers isBidirectional batchFirst))
         => ParameterNames (GRUParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) where
    type NamedParameters (GRUParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) =
      NamedParameters (GenericRNNParam ty ki 3 inF hiddenF nrLayers isBidirectional batchFirst)
    namedParameters _ = namedParameters (Proxy :: Proxy (GenericRNNParam ty ki 3 inF hiddenF nrLayers isBidirectional batchFirst))
  
  instance (ParameterNames (GenericRNNParam ty ki 4 inF hiddenF nrLayers isBidirectional batchFirst))
         => ParameterNames (LSTMParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) where
    type NamedParameters (LSTMParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) =
      NamedParameters (GenericRNNParam ty ki 4 inF hiddenF nrLayers isBidirectional batchFirst)
    namedParameters _ = namedParameters (Proxy :: Proxy (GenericRNNParam ty ki 4 inF hiddenF nrLayers isBidirectional batchFirst))
  
  instance ParameterNames (RNNCellParam ty ki inF hiddenF) where
    type NamedParameters (RNNCellParam ty ki inF hiddenF) = (Text, Text, Text, Text)
    namedParameters _ (t,t',t'',t''') = ParameterLeaf [t,t',t'',t''']
  
  instance ParameterNames (BatchNormState ty ki sz) where
    type NamedParameters (BatchNormState ty ki sz) = (Text, Text)
    namedParameters _ (t,t') = ParameterLeaf [t,t']
  
  -- | Internal. Nominally this is unsafe, the ParameterTree is assumed to always
  -- have the right shape. In reality, ParameterNames guarantees that this is
  -- ok. It would be nice to have just one integrated version, but the resulting
  -- code that does both is an unreadable mess.
  class EotStoredWithNames eot where
    eotLoadWithNames :: StoredModel -> ParameterTree -> IO eot
    eotSaveWithNames :: StoredModel -> ParameterTree -> eot -> IO StoredModel
  
  instance {-# OVERLAPS #-} (Stored x, ParameterNames x, EotStoredWithNames xs) => EotStoredWithNames (x, xs) where
    eotLoadWithNames sm (ParameterBranch (p:ps)) = do
      v <- loadWithNamesUnsafe sm p
      vs <- eotLoadWithNames sm (ParameterBranch ps)
      pure (v, vs)
    eotSaveWithNames sm (ParameterBranch (p:ps)) (x,xs) = do
      saveWithNamesUnsafe sm p x
      eotSaveWithNames sm (ParameterBranch ps) xs
      pure sm
  
  instance {-# OVERLAPS #-} (Stored x, ParameterNames x) => EotStoredWithNames (x, ()) where
    eotLoadWithNames sm (ParameterLeaf l) =
      (,()) <$> loadWithNamesUnsafe sm (ParameterLeaf l)
    eotLoadWithNames sm (ParameterBranch [p]) =
      (,()) <$> loadWithNamesUnsafe sm p
    eotSaveWithNames sm (ParameterLeaf l) (x,()) = do
      saveWithNamesUnsafe sm (ParameterLeaf l) x
      pure sm
    eotSaveWithNames sm (ParameterBranch [p]) (x,()) = do
      saveWithNamesUnsafe sm p x
      pure sm
  
  instance (EotStoredWithNames fields) => EotStoredWithNames (Either fields Void) where
    eotLoadWithNames sm names = Left <$> eotLoadWithNames sm names
    eotSaveWithNames sm names (Left x) = eotSaveWithNames sm names x
  
  instance EotStoredWithNames () where
    eotLoadWithNames _ _ = pure ()
    eotSaveWithNames sm _ _ = pure sm
  
  -- Nominally this is unsafe, the ParameterTree is assumed to always have the
  -- right shape. In reality, ParameterNames guarantees that this is ok.
  class EotStored meta eot where
    eotLoadWithPrefix :: meta -> StoredModel -> Text -> IO eot
    eotSaveWithPrefix :: meta -> StoredModel -> Text -> eot -> IO StoredModel
  
  instance (EotStored [String] fields, EotStored Int fields) =>
    EotStored GE.Datatype (Either fields Void) where
    eotLoadWithPrefix datatype sm prefix = case datatype of
      Datatype tyName [Constructor con (Selectors fieldNames)] ->
        Left <$> eotLoadWithPrefix fieldNames sm prefix
      Datatype tyName [Constructor con (NoSelectors nrFields)] ->
        Left <$> eotLoadWithPrefix (0::Int) sm prefix
      Datatype tyName _ ->
        error ("type " ++ tyName ++ " must have exactly one constructor")
    eotSaveWithPrefix datatype sm prefix x = case datatype of
      Datatype tyName [Constructor con (Selectors fieldNames)] ->
        eotSaveWithPrefix fieldNames sm prefix x
      Datatype tyName [Constructor con (NoSelectors nrFields)] ->
        eotSaveWithPrefix (0::Int) sm prefix x
      Datatype tyName _ ->
        error ("type " ++ tyName ++ " must have exactly one constructor")
  
  instance (EotStored [String] fields) => EotStored [String] (Either fields Void) where
    eotLoadWithPrefix fieldNames sm prefix = Left <$> eotLoadWithPrefix fieldNames sm prefix
    eotSaveWithPrefix fieldNames sm prefix (Left x) = eotSaveWithPrefix fieldNames sm prefix x
  
  instance (EotStored Int fields) => EotStored Int (Either fields Void) where
    eotLoadWithPrefix fieldNr sm prefix = Left <$> eotLoadWithPrefix fieldNr sm prefix
    eotSaveWithPrefix fieldNr sm prefix (Left x) = eotSaveWithPrefix fieldNr sm prefix x
  
  instance (Stored x, ParameterNames x, EotStored [String] xs) => EotStored [String] (x, xs) where
    eotLoadWithPrefix (n:ns) sm prefix = do
      v <- loadWithPrefix sm (addStoredModulePath prefix (T.pack n))
      vs <- eotLoadWithPrefix ns sm prefix
      pure (v, vs)
    eotSaveWithPrefix (n:ns) sm prefix (x,xs) = do
      saveWithPrefix sm (addStoredModulePath prefix (T.pack n)) x
      eotSaveWithPrefix ns sm prefix xs
      pure sm
  
  instance (Stored x, ParameterNames x, EotStored Int xs) => EotStored Int (x, xs) where
    eotLoadWithPrefix i sm prefix = do
      v <- loadWithPrefix sm (addStoredModulePath prefix (T.pack (show i)))
      vs <- eotLoadWithPrefix (i + 1) sm prefix
      pure (v, vs)
    eotSaveWithPrefix i sm prefix (x,xs) = do
      saveWithPrefix sm (addStoredModulePath prefix (T.pack (show i))) x
      eotSaveWithPrefix (i + 1) sm prefix xs
      pure sm
  
  instance EotStored [String] () where
    eotLoadWithPrefix _ _ _ = pure ()
    eotSaveWithPrefix _ sm _ _ = pure sm
  
  instance EotStored Int () where
    eotLoadWithPrefix _ _ _ = pure ()
    eotSaveWithPrefix _ sm _ _ = pure sm
  
  instance (TensorConstraints ty ki sz) => Stored (Tensor ty ki sz) where
    loadWithPrefix mod names =
      loadWithNamesUnsafe mod (ParameterLeaf [names])
    loadWithNames mod names =
      loadWithNamesUnsafe mod (namedParameters (Proxy @(Tensor ty ki sz)) names)
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [name]) = do
      m <- readIORef mod
      toDevice $ lookupTensorInSModule name m
    saveWithPrefix mod names ten =
      saveWithNamesUnsafe mod (ParameterLeaf [names]) ten
    saveWithNames mod names ten =
      saveWithNamesUnsafe mod (namedParameters (Proxy @(Tensor ty ki sz)) names) ten
    saveWithNamesUnsafe mod (ParameterLeaf [name]) ten =
      addTensorToStoredModel mod name =<< toCpu ten
  
  -- | Internal, used for building layers. This can't be a default because we
  -- reserve those for Generic instances, that's what users will create.
  simpleLoadWithPrefix :: (Stored a, ParameterNames a)
                       => [Text] -> StoredModel -> Text -> IO a
  simpleLoadWithPrefix params mod names = loadWithNamesUnsafe mod (ParameterLeaf (map (addStoredModulePath names) params))
  
  -- | Internal, used for building layers. This can't be a default because we
  -- reserve those for Generic instances, that's what users will create.
  simpleLoadWithNames :: forall a. (Stored a, ParameterNames a)
                      => StoredModel -> NamedParameters a -> IO a
  simpleLoadWithNames mod names = loadWithNamesUnsafe mod (namedParameters (Proxy @a) names)
  
  -- | Internal, used for building layers. This can't be a default because we
  -- reserve those for Generic instances, that's what users will create.
  simpleSaveWithPrefix :: (Stored a, ParameterNames a)
                       => [Text] -> StoredModel -> Text -> a -> IO StoredModel
  simpleSaveWithPrefix params mod names ten = saveWithNamesUnsafe mod (ParameterLeaf (map (addStoredModulePath names) params)) ten
  
  -- | Internal, used for building layers. This can't be a default because we
  -- reserve those for Generic instances, that's what users will create.
  simpleSaveWithNames :: forall a. (Stored a, ParameterNames a)
                      => StoredModel -> NamedParameters a -> a -> IO StoredModel
  simpleSaveWithNames mod names ten = saveWithNamesUnsafe mod (namedParameters (Proxy @a) names) ten
  
  instance (TensorConstraints ty ki inF, SingI inF, SingI outF) => Stored (LinearParam ty ki inF outF) where
    loadWithPrefix = simpleLoadWithPrefix ["weights", "biases"]
    saveWithPrefix = simpleSaveWithPrefix ["weights", "biases"]
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [wname, bname]) = do
      m <- readIORef mod
      LinearParam <$> toDevice (lookupTensorInSModule wname m)
                  <*> mapM toDevice (maybeLookupTensorInSModule' bname m)
    saveWithNamesUnsafe mod (ParameterLeaf [wname, bname]) (LinearParam wt Nothing) =
      addTensorToStoredModel mod wname =<< toCpu wt
    saveWithNamesUnsafe mod (ParameterLeaf [wname, bname]) (LinearParam wt (Just bt)) = do
      mod' <- addTensorToStoredModel mod wname =<< toCpu wt
      addTensorToStoredModel mod' bname =<< toCpu bt
  
  instance (TensorConstraints ty ki outChans, SingI outChans, SingI otherSz) => Stored (ConvParam ty ki outChans otherSz) where
    loadWithPrefix = simpleLoadWithPrefix ["weights", "biases"]
    saveWithPrefix = simpleSaveWithPrefix ["weights", "biases"]
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [wname, bname]) = do
      m <- readIORef mod
      ConvParam <$> toDevice (lookupTensorInSModule wname m)
                <*> mapM toDevice (maybeLookupTensorInSModule' bname m)
    saveWithNamesUnsafe mod (ParameterLeaf [wname, bname]) (ConvParam wt Nothing) =
      addTensorToStoredModel mod wname =<< toCpu wt
    saveWithNamesUnsafe mod (ParameterLeaf [wname, bname]) (ConvParam wt (Just bt)) = do
      mod' <- addTensorToStoredModel mod wname =<< toCpu wt
      addTensorToStoredModel mod' bname =<< toCpu bt
  
  instance (TensorConstraints ty ki sz) => Stored (AffineParam ty ki sz) where
    loadWithPrefix = simpleLoadWithPrefix ["weights", "biases"]
    saveWithPrefix = simpleSaveWithPrefix ["weights", "biases"]
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [wname, bname]) = do
      m <- readIORef mod
      AffineParam <$> toDevice (lookupTensorInSModule wname m)
                  <*> mapM toDevice (maybeLookupTensorInSModule' bname m)
    saveWithNamesUnsafe mod (ParameterLeaf [wname, bname]) (AffineParam wt Nothing) =
      addTensorToStoredModel mod wname =<< toCpu wt
    saveWithNamesUnsafe mod (ParameterLeaf [wname, bname]) (AffineParam wt (Just bt)) = do
      mod' <- addTensorToStoredModel mod wname =<< toCpu wt
      addTensorToStoredModel mod' bname =<< toCpu bt
  
  instance (TensorConstraints ty ki '[nrEmbeddings, embeddingDim], SingI nrEmbeddings, SingI embeddingDim)
         => Stored (EmbeddingParam ty ki nrEmbeddings embeddingDim) where
    loadWithPrefix = simpleLoadWithPrefix ["embedding"]
    saveWithPrefix = simpleSaveWithPrefix ["embedding"]
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [ename]) = do
      m <- readIORef mod
      EmbeddingParam <$> toDevice (lookupTensorInSModule ename m)
    saveWithNamesUnsafe mod (ParameterLeaf [ename]) (EmbeddingParam embed) =
      addTensorToStoredModel mod ename =<< toCpu embed
  
  instance (TensorConstraints ty ki sz) => Stored (BatchNormState ty ki sz) where
    loadWithPrefix = simpleLoadWithPrefix ["means", "variances"]
    saveWithPrefix = simpleSaveWithPrefix ["means", "variances"]
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [mname, vname]) = do
      m <- readIORef mod
      mten <- mapM toDevice $ maybeLookupTensorInSModule' mname m
      vten <- mapM toDevice $ maybeLookupTensorInSModule' vname m
      BatchNormState <$> newIORef (BatchNormData mten vten)
    saveWithNamesUnsafe mod (ParameterLeaf [mname, vname]) (BatchNormState bn) = do
      (BatchNormData m v) <- readIORef bn
      mod' <- case m of
               Nothing -> pure mod
               Just mt -> addTensorToStoredModel mod mname =<< toCpu mt
      case v of
        Nothing -> pure mod'
        Just vt -> addTensorToStoredModel mod' vname =<< toCpu vt
  
  instance (TensorConstraints ty ki inF, SingI inF, SingI hiddenF, SingI nrLayers, SingI (NrOfRNNDirections isBidirectional)
           ,SingI isBidirectional, SingI (NrOfRNNDirections isBidirectional TL.* hiddenF)
           ,SingI (gateSize TL.* hiddenF))
          => Stored (GenericRNNParam ty ki gateSize inF hiddenF nrLayers isBidirectional batchFirst) where
    loadWithPrefix =
      let nrOfTensors = demoteN @nrLayers * demoteN @(NrOfRNNDirections isBidirectional)
          layerNames label = map (\n -> label <> T.pack (show n)) [0..nrOfTensors-1]
      in simpleLoadWithPrefix (layerNames "wih" ++ layerNames "bih" ++ layerNames "whh" ++ layerNames "bhh")
    saveWithPrefix =
      let nrOfTensors = demoteN @nrLayers * demoteN @(NrOfRNNDirections isBidirectional)
          layerNames label = map (\n -> label <> T.pack (show n)) [0..nrOfTensors-1]
      in simpleSaveWithPrefix (layerNames "wih" ++ layerNames "bih" ++ layerNames "whh" ++ layerNames "bhh")
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf names) = do
      let nrDirections = demoteN @(NrOfRNNDirections isBidirectional)
      let nrOfTensors = demoteN @nrLayers * nrDirections
      let loadList m l = fromListMDirection (toDevice . flip lookupTensorInSModule m) l
      m <- readIORef mod
      let [namesWih, namesBih, namesWhh, namesBhh] = splitEvery nrOfTensors names
      let hasBias = isJust $ M.lookup (head namesBih) (smodelTensors m)
      wih0  <- loadList m (Prelude.take nrDirections namesWih)
      wihN1 <- VB.fromList <$> mapM (loadList m) (splitEvery nrDirections $ drop nrDirections namesWih)
      bihN  <- if hasBias then
                Just <$> VB.fromList <$> mapM (loadList m) (splitEvery nrDirections namesBih)
                else pure Nothing
      whhN  <- VB.fromList <$> mapM (loadList m) (splitEvery nrDirections namesWhh)
      bhhN  <- if hasBias then
                Just <$> VB.fromList <$> mapM (loadList m) (splitEvery nrDirections namesBhh)
                else pure Nothing
      pure $ GenericRNNParam wih0 wihN1 bihN whhN bhhN
    saveWithNamesUnsafe mod (ParameterLeaf names) (GenericRNNParam wih0 wihN1 maybeBihN whhN maybeBhhN) = do
      let nrDirections = demoteN @(NrOfRNNDirections isBidirectional)
      let nrOfTensors = demoteN @nrLayers * nrDirections
      let [namesWih, namesBih, namesWhh, namesBhh] = splitEvery nrOfTensors names
      let saveList :: (SingI sz) => PerDirection isBidirectional (Tensor ty ki sz) -> [Text] -> IO (PerDirection isBidirectional StoredModel)
          saveList tens names = zipWithMDirection (\t n -> addTensorToStoredModel mod n =<< toCpu t) tens names
      saveList wih0 (Prelude.take nrDirections namesWih)
      zipWithM_ saveList (VB.toList wihN1) (splitEvery nrDirections $ drop nrDirections namesWih)
      case maybeBihN of
        Just bihN -> zipWithM_ saveList (VB.toList bihN)  (splitEvery nrDirections namesBih)
        Nothing   -> pure ()
      zipWithM_ saveList (VB.toList whhN) (splitEvery nrDirections namesWhh)
      case maybeBhhN of
        Just bhhN -> zipWithM_ saveList (VB.toList bhhN)  (splitEvery nrDirections namesBhh)
        Nothing   -> pure ()
      pure mod
  
  instance (TensorConstraints ty ki inF, SingI inF, SingI hiddenF
           ,SingI nrLayers, SingI (NrOfRNNDirections isBidirectional)
           ,Stored (GenericRNNParam ty ki 1 inF hiddenF nrLayers isBidirectional batchFirst))
         => Stored (RNNParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) where
    loadWithPrefix sm t = RNNParams <$> loadWithPrefix sm t
    saveWithPrefix sm t (RNNParams rnn) = saveWithPrefix sm t rnn
    loadWithNames sm ns = RNNParams <$> loadWithNames sm ns
    saveWithNames sm ns (RNNParams rnn) = saveWithNames sm ns rnn
    loadWithNamesUnsafe sm ns = RNNParams <$> loadWithNamesUnsafe sm ns
    saveWithNamesUnsafe sm ns (RNNParams rnn) = saveWithNamesUnsafe sm ns rnn
  
  instance (TensorConstraints ty ki inF, SingI inF, SingI hiddenF
           ,SingI isBidirectional, SingI (NrOfRNNDirections isBidirectional TL.* hiddenF)
           ,SingI nrLayers, SingI (NrOfRNNDirections isBidirectional)
           ,Stored (GenericRNNParam ty ki 3 inF hiddenF nrLayers isBidirectional batchFirst))
         => Stored (GRUParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) where
    loadWithPrefix sm t = GRUParams <$> loadWithPrefix sm t
    saveWithPrefix sm t (GRUParams rnn) = saveWithPrefix sm t rnn
    loadWithNames sm ns = GRUParams <$> loadWithNames sm ns
    saveWithNames sm ns (GRUParams rnn) = saveWithNames sm ns rnn
    loadWithNamesUnsafe sm ns = GRUParams <$> loadWithNamesUnsafe sm ns
    saveWithNamesUnsafe sm ns (GRUParams rnn) = saveWithNamesUnsafe sm ns rnn
  
  instance (TensorConstraints ty ki inF, SingI inF, SingI hiddenF
           ,SingI isBidirectional, SingI (NrOfRNNDirections isBidirectional TL.* hiddenF)
           ,SingI nrLayers, SingI (NrOfRNNDirections isBidirectional)
           ,Stored (GenericRNNParam ty ki 4 inF hiddenF nrLayers isBidirectional batchFirst))
         => Stored (LSTMParams ty ki inF hiddenF nrLayers isBidirectional batchFirst) where
    loadWithPrefix sm t = LSTMParams <$> loadWithPrefix sm t
    saveWithPrefix sm t (LSTMParams rnn) = saveWithPrefix sm t rnn
    loadWithNames sm ns = LSTMParams <$> loadWithNames sm ns
    saveWithNames sm ns (LSTMParams rnn) = saveWithNames sm ns rnn
    loadWithNamesUnsafe sm ns = LSTMParams <$> loadWithNamesUnsafe sm ns
    saveWithNamesUnsafe sm ns (LSTMParams rnn) = saveWithNamesUnsafe sm ns rnn
  
  instance (TensorConstraints ty ki inF, SingI inF, SingI hiddenF) => Stored (RNNCellParam ty ki inF hiddenF) where
    loadWithPrefix = simpleLoadWithPrefix ["wih", "bih", "whh", "bhh"]
    saveWithPrefix = simpleSaveWithPrefix ["wih", "bih", "whh", "bhh"]
    loadWithNames = simpleLoadWithNames
    saveWithNames = simpleSaveWithNames
    loadWithNamesUnsafe (StoredModel mod) (ParameterLeaf [wihname, bihname, whhname, bhhname, statename]) = do
      m <- readIORef mod
      RNNCellParam <$> toDevice (lookupTensorInSModule wihname m)
                   <*> mapM toDevice (maybeLookupTensorInSModule' bihname m)
                   <*> toDevice (lookupTensorInSModule whhname m)
                   <*> mapM toDevice (maybeLookupTensorInSModule' bhhname m)
    saveWithNamesUnsafe mod (ParameterLeaf [wihname, bihname, whhname, bhhname])
                            (RNNCellParam wiht Nothing whht Nothing) = do
      addTensorToStoredModel mod wihname =<< toCpu wiht
      addTensorToStoredModel mod whhname =<< toCpu whht
    saveWithNamesUnsafe mod (ParameterLeaf [wihname, bihname, whhname, bhhname])
                            (RNNCellParam wiht (Just biht) whht (Just bhht)) = do
      addTensorToStoredModel mod wihname =<< toCpu wiht
      addTensorToStoredModel mod bihname =<< toCpu biht
      addTensorToStoredModel mod whhname =<< toCpu whht
      addTensorToStoredModel mod bhhname =<< toCpu bhht
  
  -- * Quick convenience functions
  
  writeModelToFile :: Stored a => a -> Text -> IO ()
  writeModelToFile ten fname = do
    sm <- emptyStoredModel
    save sm ten
    writeStoredModel sm fname
  
  readModelFromFile :: Stored b => Text -> IO b
  readModelFromFile fname = do
    sm <- readStoredModel fname
    load sm
  
  writeModelToFileWithPrefix :: Stored a => a -> Text -> Text -> IO ()
  writeModelToFileWithPrefix ten prefix fname = do
    sm <- emptyStoredModel
    saveWithPrefix sm prefix ten
    writeStoredModel sm fname
  
  readModelFromFileWithPrefix :: Stored b => Text -> Text -> IO b
  readModelFromFileWithPrefix prefix fname = do
    sm <- readStoredModel fname
    loadWithPrefix sm prefix
  
  writeModelToFileWithNames :: (Stored a, ParameterNames a) => a -> NamedParameters a -> Text -> IO ()
  writeModelToFileWithNames ten names fname = do
    sm <- emptyStoredModel
    saveWithNames sm names ten
    writeStoredModel sm fname
  
  readModelFromFileWithNames :: (Stored b, ParameterNames b) => NamedParameters b -> Text -> IO b
  readModelFromFileWithNames names fname = do
    sm <- readStoredModel fname
    loadWithNames sm names
  
  -- * Working with stored models, these are lower-level operations
  
  -- | Create a new model
  emptyStoredModel :: IO StoredModel
  emptyStoredModel = do
    m <- C.mkJitModule
    StoredModel <$> newIORef (StoredModelData { smodelRoot     = m
                                              , smodelEntities = M.empty
                                              , smodelTensors  = M.empty
                                              })
  
  -- | Read a stored model from disk
  readStoredModel :: Text -> IO StoredModel
  readStoredModel str = do
    m <- C.jitModuleReadFromFile str
    mod <- case m of
      Nothing  -> error $ T.unpack $ "Can't open file containing stored module " <> str
      Just mod -> pure mod
    -- TODO When CUDA is available models are loaded into it directly. That would
    -- be fine, but we need to make StoredModel contain the type in this
    -- case. Temporarily we just make sure that everything is on the CPU.
    C.jitModueToBackend mod 0
    c <- storedModelContents mod
    tensorMap <- convertStoredModelTensors mod
    StoredModel <$> newIORef (StoredModelData { smodelRoot     = mod
                                              , smodelEntities = M.fromList c
                                              , smodelTensors  = tensorMap
                                              })
  
  -- | Write a stored model to disk
  writeStoredModel :: StoredModel -> Text -> IO ()
  writeStoredModel (StoredModel sm) fname = do
    smd <- readIORef sm
    C.jitModueWriteToFile (smodelRoot smd) fname
  
  -- | Models are like dictionarie. Add a tensor to a model.
  addTensorToStoredModel :: StoredModel -> Text -> Tensor ty KCpu sz -> IO StoredModel
  addTensorToStoredModel s@(StoredModel sm) key (Tensor ten _) = do
    smd <- readIORef sm
    C.jitModuleAddTensor (smodelRoot smd) key ten
    modifyIORef' sm (\sm -> sm { smodelEntities = M.insert key C.ParameterEntityType $ smodelEntities sm
                              , smodelTensors  = M.insert key ten $ smodelTensors sm
                              })
    pure s
  
  -- | List the names of the tensors in a model
  listStoredModelTensors :: StoredModel -> IO [Text]
  listStoredModelTensors (StoredModel sm) = do
    smd <- readIORef sm
    pure $ M.keys $ smodelTensors smd
  
  -- * Internal functions, helpful if you make your own layers
  
  -- Internal
  storedModelContents :: ForeignPtr C.CJitScriptModule -> IO [(Text, C.ModuleEntityType)]
  storedModelContents mod = do
    nr <- C.jitModuleNumberOfSlots mod
    mapM (\i -> do
             s <- C.jitModuleSlotName mod (fromIntegral i)
             isModule <- C.jitModuleSlotIsModule mod (fromIntegral i)
             isParameter <- C.jitModuleSlotIsParameter mod (fromIntegral i)
             pure (s,case (isModule, isParameter) of
                      (1, 0) -> C.ModuleEntityType
                      (0, 1) -> C.ParameterEntityType
                      (0, 0) -> C.MethodEntityType)) [0..(nr-1)]
  
  -- Internal
  convertStoredModelTensors :: ForeignPtr C.CJitScriptModule -> IO (M.Map Text (ForeignPtr C.CTensor))
  convertStoredModelTensors mod = do
    cs <- storedModelContents mod
    M.fromList <$> concat <$> mapM (process mod "") cs
    where process mod path (name, ty) =
             case ty of
               C.ModuleEntityType -> do
                 mod' <- C.jitModuleReadModule mod name
                 cs <- storedModelContents mod'
                 concat <$> mapM (process mod' (addPath path name)) cs
               C.ParameterEntityType -> do
                 ten <- C.jitModuleReadTensor mod name
                 pure [(addPath path name, ten)]
          addPath "" name = name
          addPath p  name = p <> "." <> name
  
  -- Internal
  lookupTensorInSModule :: TensorConstraints ty KCpu sz => Text -> StoredModelData -> Tensor ty KCpu sz
  lookupTensorInSModule name mod =
    case M.lookup name (smodelTensors mod) of
      Nothing -> error $ "Can't find required entry in serialized module " <> T.unpack name
      -- TODO Does this tensor need to hold on to the module?
      Just t  -> debuggingVerifyShape "Looking up stored tensor" $ Tensor t Nothing
  
  -- Internal
  maybeLookupTensorInSModule :: TensorConstraints ty KCpu sz
                             => Maybe Text -> StoredModelData -> Maybe (Tensor ty KCpu sz)
  maybeLookupTensorInSModule Nothing mod     = Nothing
  maybeLookupTensorInSModule (Just name) mod = Just $ lookupTensorInSModule name mod
  
  -- Internal
  maybeLookupTensorInSModule' :: TensorConstraints ty KCpu sz
                              => Text -> StoredModelData -> Maybe (Tensor ty KCpu sz)
  maybeLookupTensorInSModule' name mod =
    case M.lookup name (smodelTensors mod) of
      Nothing -> Nothing
      Just t  -> Just $ debuggingVerifyShape "Looking up optional stored tensor" $ Tensor t Nothing
  
  addStoredModulePath "" next     = next
  addStoredModulePath prefix next = prefix <> "." <> next
  