{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UndecidableSuperClasses #-}
{-# LANGUAGE NoStarIsType #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Extra.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

module Torch.Typed.NN.Recurrent.GRU where

import Data.Kind
import Data.Proxy (Proxy (..))
import Foreign.ForeignPtr
import GHC.Generics
import GHC.TypeLits
import GHC.TypeLits.Extra
import System.Environment
import System.IO.Unsafe
import qualified Torch.Autograd as A
import qualified Torch.DType as D
import qualified Torch.Device as D
import qualified Torch.Functional as D
import Torch.HList
import qualified Torch.Internal.Cast as ATen
import qualified Torch.Internal.Class as ATen
import qualified Torch.Internal.Managed.Type.Tensor as ATen
import qualified Torch.Internal.Type as ATen
import qualified Torch.NN as A
import qualified Torch.Tensor as D
import qualified Torch.TensorFactories as D
import Torch.Typed.Factories
import Torch.Typed.Functional hiding (sqrt)
import Torch.Typed.NN.Dropout
import Torch.Typed.NN.Recurrent.Auxiliary
import Torch.Typed.Parameter
import Torch.Typed.Tensor
import Prelude hiding (tanh)

data
  GRULayerSpec
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (directionality :: RNNDirectionality)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  = GRULayerSpec
  deriving (Int
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Int
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
[GRULayerSpec inputSize hiddenSize directionality dtype device]
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
-> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GRULayerSpec inputSize hiddenSize directionality dtype device]
-> ShowS
$cshowList :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
[GRULayerSpec inputSize hiddenSize directionality dtype device]
-> ShowS
show :: GRULayerSpec inputSize hiddenSize directionality dtype device
-> String
$cshow :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
-> String
showsPrec :: Int
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> ShowS
$cshowsPrec :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Int
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> ShowS
Show, GRULayerSpec inputSize hiddenSize directionality dtype device
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> Bool
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GRULayerSpec inputSize hiddenSize directionality dtype device
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> Bool
$c/= :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> Bool
== :: GRULayerSpec inputSize hiddenSize directionality dtype device
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> Bool
$c== :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
-> GRULayerSpec inputSize hiddenSize directionality dtype device
-> Bool
Eq)

data
  GRULayer
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (directionality :: RNNDirectionality)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  where
  GRUUnidirectionalLayer ::
    Parameter device dtype (GRUWIShape hiddenSize inputSize) ->
    Parameter device dtype (GRUWHShape hiddenSize inputSize) ->
    Parameter device dtype (GRUBIShape hiddenSize inputSize) ->
    Parameter device dtype (GRUBHShape hiddenSize inputSize) ->
    GRULayer inputSize hiddenSize 'Unidirectional dtype device
  GRUBidirectionalLayer ::
    Parameter device dtype (GRUWIShape hiddenSize inputSize) ->
    Parameter device dtype (GRUWHShape hiddenSize inputSize) ->
    Parameter device dtype (GRUBIShape hiddenSize inputSize) ->
    Parameter device dtype (GRUBHShape hiddenSize inputSize) ->
    Parameter device dtype (GRUWIShape hiddenSize inputSize) ->
    Parameter device dtype (GRUWHShape hiddenSize inputSize) ->
    Parameter device dtype (GRUBIShape hiddenSize inputSize) ->
    Parameter device dtype (GRUBHShape hiddenSize inputSize) ->
    GRULayer inputSize hiddenSize 'Bidirectional dtype device

deriving instance Show (GRULayer inputSize hiddenSize directionality dtype device)

instance Parameterized (GRULayer inputSize hiddenSize 'Unidirectional dtype device) where
  type
    Parameters (GRULayer inputSize hiddenSize 'Unidirectional dtype device) =
      '[ Parameter device dtype (GRUWIShape hiddenSize inputSize),
         Parameter device dtype (GRUWHShape hiddenSize inputSize),
         Parameter device dtype (GRUBIShape hiddenSize inputSize),
         Parameter device dtype (GRUBHShape hiddenSize inputSize)
       ]
  flattenParameters :: GRULayer inputSize hiddenSize 'Unidirectional dtype device
-> HList
     (Parameters
        (GRULayer inputSize hiddenSize 'Unidirectional dtype device))
flattenParameters (GRUUnidirectionalLayer Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh) =
    Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. forall k. HList '[]
HNil
  replaceParameters :: GRULayer inputSize hiddenSize 'Unidirectional dtype device
-> HList
     (Parameters
        (GRULayer inputSize hiddenSize 'Unidirectional dtype device))
-> GRULayer inputSize hiddenSize 'Unidirectional dtype device
replaceParameters GRULayer inputSize hiddenSize 'Unidirectional dtype device
_ (Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi :. Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh :. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi :. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh :. HList '[]
R:HListk[] Type
HNil) =
    forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (inputSize :: Natural).
Parameter device dtype (GRUWIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWHShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> GRULayer inputSize hiddenSize 'Unidirectional dtype device
GRUUnidirectionalLayer Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh

instance Parameterized (GRULayer inputSize hiddenSize 'Bidirectional dtype device) where
  type
    Parameters (GRULayer inputSize hiddenSize 'Bidirectional dtype device) =
      '[ Parameter device dtype (GRUWIShape hiddenSize inputSize),
         Parameter device dtype (GRUWHShape hiddenSize inputSize),
         Parameter device dtype (GRUBIShape hiddenSize inputSize),
         Parameter device dtype (GRUBHShape hiddenSize inputSize),
         Parameter device dtype (GRUWIShape hiddenSize inputSize),
         Parameter device dtype (GRUWHShape hiddenSize inputSize),
         Parameter device dtype (GRUBIShape hiddenSize inputSize),
         Parameter device dtype (GRUBHShape hiddenSize inputSize)
       ]
  flattenParameters :: GRULayer inputSize hiddenSize 'Bidirectional dtype device
-> HList
     (Parameters
        (GRULayer inputSize hiddenSize 'Bidirectional dtype device))
flattenParameters (GRUBidirectionalLayer Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi' Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh' Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi' Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh') =
    Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi' forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh' forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi' forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh' forall x (xs :: [Type]). x -> HList xs -> HList (x : xs)
:. forall k. HList '[]
HNil
  replaceParameters :: GRULayer inputSize hiddenSize 'Bidirectional dtype device
-> HList
     (Parameters
        (GRULayer inputSize hiddenSize 'Bidirectional dtype device))
-> GRULayer inputSize hiddenSize 'Bidirectional dtype device
replaceParameters GRULayer inputSize hiddenSize 'Bidirectional dtype device
_ (Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi :. Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh :. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi :. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh :. Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi' :. Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh' :. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi' :. Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh' :. HList '[]
R:HListk[] Type
HNil) =
    forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (inputSize :: Natural).
Parameter device dtype (GRUWIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWHShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWHShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> GRULayer inputSize hiddenSize 'Bidirectional dtype device
GRUBidirectionalLayer Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh Parameter device dtype (GRUWIShape hiddenSize inputSize)
wi' Parameter device dtype (GRUWHShape hiddenSize inputSize)
wh' Parameter device dtype (GRUBIShape hiddenSize inputSize)
bi' Parameter device dtype (GRUBIShape hiddenSize inputSize)
bh'

instance
  ( RandDTypeIsValid device dtype,
    KnownNat inputSize,
    KnownNat hiddenSize,
    KnownDType dtype,
    KnownDevice device
  ) =>
  A.Randomizable
    (GRULayerSpec inputSize hiddenSize 'Unidirectional dtype device)
    (GRULayer inputSize hiddenSize 'Unidirectional dtype device)
  where
  sample :: GRULayerSpec inputSize hiddenSize 'Unidirectional dtype device
-> IO (GRULayer inputSize hiddenSize 'Unidirectional dtype device)
sample GRULayerSpec inputSize hiddenSize 'Unidirectional dtype device
_ =
    forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (inputSize :: Natural).
Parameter device dtype (GRUWIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWHShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> GRULayer inputSize hiddenSize 'Unidirectional dtype device
GRUUnidirectionalLayer
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros)

instance
  ( RandDTypeIsValid device dtype,
    KnownNat inputSize,
    KnownNat hiddenSize,
    KnownDType dtype,
    KnownDevice device
  ) =>
  A.Randomizable
    (GRULayerSpec inputSize hiddenSize 'Bidirectional dtype device)
    (GRULayer inputSize hiddenSize 'Bidirectional dtype device)
  where
  sample :: GRULayerSpec inputSize hiddenSize 'Bidirectional dtype device
-> IO (GRULayer inputSize hiddenSize 'Bidirectional dtype device)
sample GRULayerSpec inputSize hiddenSize 'Bidirectional dtype device
_ =
    forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (inputSize :: Natural).
Parameter device dtype (GRUWIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWHShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWIShape hiddenSize inputSize)
-> Parameter device dtype (GRUWHShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> Parameter device dtype (GRUBIShape hiddenSize inputSize)
-> GRULayer inputSize hiddenSize 'Bidirectional dtype device
GRUBidirectionalLayer
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros)

data
  GRULayerStackSpec
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (numLayers :: Nat)
    (directionality :: RNNDirectionality)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  = GRULayerStackSpec
  deriving (Int
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
Int
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
[GRULayerStackSpec
   inputSize hiddenSize numLayers directionality dtype device]
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GRULayerStackSpec
   inputSize hiddenSize numLayers directionality dtype device]
-> ShowS
$cshowList :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
[GRULayerStackSpec
   inputSize hiddenSize numLayers directionality dtype device]
-> ShowS
show :: GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> String
$cshow :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> String
showsPrec :: Int
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
$cshowsPrec :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
Int
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
Show, GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> Bool
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> Bool
$c/= :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> Bool
== :: GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> Bool
$c== :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> Bool
Eq)

-- Input-to-hidden, hidden-to-hidden, and bias parameters for a mulilayered
-- (and optionally) bidirectional GRU.
--
data
  GRULayerStack
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (numLayers :: Nat)
    (directionality :: RNNDirectionality)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  where
  GRULayer1 ::
    GRULayer inputSize hiddenSize directionality dtype device ->
    GRULayerStack inputSize hiddenSize 1 directionality dtype device
  GRULayerK ::
    GRULayer (hiddenSize * NumberOfDirections directionality) hiddenSize directionality dtype device ->
    GRULayerStack inputSize hiddenSize numLayers directionality dtype device ->
    GRULayerStack inputSize hiddenSize (numLayers + 1) directionality dtype device

deriving instance Show (GRULayerStack inputSize hiddenSize numLayers directionality dtype device)

class GRULayerStackParameterized (flag :: Bool) inputSize hiddenSize numLayers directionality dtype device where
  type GRULayerStackParameters flag inputSize hiddenSize numLayers directionality dtype device :: [Type]
  gruLayerStackFlattenParameters ::
    Proxy flag ->
    GRULayerStack inputSize hiddenSize numLayers directionality dtype device ->
    HList (GRULayerStackParameters flag inputSize hiddenSize numLayers directionality dtype device)
  gruLayerStackReplaceParameters ::
    Proxy flag ->
    GRULayerStack inputSize hiddenSize numLayers directionality dtype device ->
    HList (GRULayerStackParameters flag inputSize hiddenSize numLayers directionality dtype device) ->
    GRULayerStack inputSize hiddenSize numLayers directionality dtype device

instance
  Parameterized (GRULayer inputSize hiddenSize directionality dtype device) =>
  GRULayerStackParameterized 'False inputSize hiddenSize 1 directionality dtype device
  where
  type
    GRULayerStackParameters 'False inputSize hiddenSize 1 directionality dtype device =
      Parameters (GRULayer inputSize hiddenSize directionality dtype device)
  gruLayerStackFlattenParameters :: Proxy 'False
-> GRULayerStack inputSize hiddenSize 1 directionality dtype device
-> HList
     (GRULayerStackParameters
        'False inputSize hiddenSize 1 directionality dtype device)
gruLayerStackFlattenParameters Proxy 'False
_ (GRULayer1 GRULayer inputSize hiddenSize directionality dtype device
gruLayer) = forall f. Parameterized f => f -> HList (Parameters f)
flattenParameters GRULayer inputSize hiddenSize directionality dtype device
gruLayer
  gruLayerStackReplaceParameters :: Proxy 'False
-> GRULayerStack inputSize hiddenSize 1 directionality dtype device
-> HList
     (GRULayerStackParameters
        'False inputSize hiddenSize 1 directionality dtype device)
-> GRULayerStack inputSize hiddenSize 1 directionality dtype device
gruLayerStackReplaceParameters Proxy 'False
_ (GRULayer1 GRULayer inputSize hiddenSize directionality dtype device
gruLayer) HList
  (GRULayerStackParameters
     'False inputSize hiddenSize 1 directionality dtype device)
parameters = forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayer inputSize hiddenSize directionality dtype device
-> GRULayerStack inputSize hiddenSize 1 directionality dtype device
GRULayer1 forall a b. (a -> b) -> a -> b
$ forall f. Parameterized f => f -> HList (Parameters f) -> f
replaceParameters GRULayer inputSize hiddenSize directionality dtype device
gruLayer HList
  (GRULayerStackParameters
     'False inputSize hiddenSize 1 directionality dtype device)
parameters

instance
  ( Parameterized
      ( GRULayer
          (hiddenSize * NumberOfDirections directionality)
          hiddenSize
          directionality
          dtype
          device
      ),
    Parameterized (GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device),
    HAppendFD
      (Parameters (GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device))
      (Parameters (GRULayer (hiddenSize * NumberOfDirections directionality) hiddenSize directionality dtype device))
      (Parameters (GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device) ++ Parameters (GRULayer (hiddenSize * NumberOfDirections directionality) hiddenSize directionality dtype device))
  ) =>
  GRULayerStackParameterized 'True inputSize hiddenSize numLayers directionality dtype device
  where
  type
    GRULayerStackParameters 'True inputSize hiddenSize numLayers directionality dtype device =
      Parameters (GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device)
        ++ Parameters (GRULayer (hiddenSize * NumberOfDirections directionality) hiddenSize directionality dtype device)
  gruLayerStackFlattenParameters :: Proxy 'True
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
-> HList
     (GRULayerStackParameters
        'True inputSize hiddenSize numLayers directionality dtype device)
gruLayerStackFlattenParameters Proxy 'True
_ (GRULayerK GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
gruLayer GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gruLayerStack) =
    let parameters :: HList
  (Parameters
     (GRULayer
        (hiddenSize * NumberOfDirections directionality)
        hiddenSize
        directionality
        dtype
        device))
parameters = forall f. Parameterized f => f -> HList (Parameters f)
flattenParameters GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
gruLayer
        parameters' :: HList
  (Parameters
     (GRULayerStack
        inputSize hiddenSize (numLayers - 1) directionality dtype device))
parameters' = forall f. Parameterized f => f -> HList (Parameters f)
flattenParameters @(GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device) GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gruLayerStack
     in HList
  (Parameters
     (GRULayerStack
        inputSize hiddenSize (numLayers - 1) directionality dtype device))
parameters' forall k (a :: [k]) (b :: [k]) (ab :: [k]).
HAppendFD a b ab =>
HList a -> HList b -> HList ab
`happendFD` HList
  (Parameters
     (GRULayer
        (hiddenSize * NumberOfDirections directionality)
        hiddenSize
        directionality
        dtype
        device))
parameters
  gruLayerStackReplaceParameters :: Proxy 'True
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
-> HList
     (GRULayerStackParameters
        'True inputSize hiddenSize numLayers directionality dtype device)
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
gruLayerStackReplaceParameters Proxy 'True
_ (GRULayerK GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
gruLayer GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gruLayerStack) HList
  (GRULayerStackParameters
     'True inputSize hiddenSize numLayers directionality dtype device)
parameters'' =
    let (HList
  (GRULayerStackParameters
     (OrdCond (CmpNat 2 (numLayers - 1)) 'True 'True 'False)
     inputSize
     hiddenSize
     (numLayers - 1)
     directionality
     dtype
     device)
parameters', HList
  (Parameters
     (GRULayer
        (hiddenSize * NumberOfDirections directionality)
        hiddenSize
        directionality
        dtype
        device))
parameters) = forall k (a :: [k]) (b :: [k]) (ab :: [k]).
HAppendFD a b ab =>
HList ab -> (HList a, HList b)
hunappendFD HList
  (GRULayerStackParameters
     'True inputSize hiddenSize numLayers directionality dtype device)
parameters''
        gruLayer' :: GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
gruLayer' = forall f. Parameterized f => f -> HList (Parameters f) -> f
replaceParameters GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
gruLayer HList
  (Parameters
     (GRULayer
        (hiddenSize * NumberOfDirections directionality)
        hiddenSize
        directionality
        dtype
        device))
parameters
        gruLayerStack' :: GRULayerStack
  inputSize hiddenSize (numLayers - 1) directionality dtype device
gruLayerStack' = forall f. Parameterized f => f -> HList (Parameters f) -> f
replaceParameters @(GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device) GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gruLayerStack HList
  (GRULayerStackParameters
     (OrdCond (CmpNat 2 (numLayers - 1)) 'True 'True 'False)
     inputSize
     hiddenSize
     (numLayers - 1)
     directionality
     dtype
     device)
parameters'
     in forall (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)) (inputSize :: Natural)
       (numLayers :: Natural).
GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStack
     inputSize hiddenSize (numLayers + 1) directionality dtype device
GRULayerK GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
gruLayer' GRULayerStack
  inputSize hiddenSize (numLayers - 1) directionality dtype device
gruLayerStack'

instance
  ( 1 <= numLayers,
    (2 <=? numLayers) ~ flag,
    GRULayerStackParameterized flag inputSize hiddenSize numLayers directionality dtype device
  ) =>
  Parameterized (GRULayerStack inputSize hiddenSize numLayers directionality dtype device)
  where
  type
    Parameters (GRULayerStack inputSize hiddenSize numLayers directionality dtype device) =
      GRULayerStackParameters (2 <=? numLayers) inputSize hiddenSize numLayers directionality dtype device
  flattenParameters :: GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
-> HList
     (Parameters
        (GRULayerStack
           inputSize hiddenSize numLayers directionality dtype device))
flattenParameters = forall (flag :: Bool) (inputSize :: Natural)
       (hiddenSize :: Natural) (numLayers :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerStackParameterized
  flag inputSize hiddenSize numLayers directionality dtype device =>
Proxy flag
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
-> HList
     (GRULayerStackParameters
        flag inputSize hiddenSize numLayers directionality dtype device)
gruLayerStackFlattenParameters (forall {k} (t :: k). Proxy t
Proxy :: Proxy flag)
  replaceParameters :: GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
-> HList
     (Parameters
        (GRULayerStack
           inputSize hiddenSize numLayers directionality dtype device))
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
replaceParameters = forall (flag :: Bool) (inputSize :: Natural)
       (hiddenSize :: Natural) (numLayers :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerStackParameterized
  flag inputSize hiddenSize numLayers directionality dtype device =>
Proxy flag
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
-> HList
     (GRULayerStackParameters
        flag inputSize hiddenSize numLayers directionality dtype device)
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
gruLayerStackReplaceParameters (forall {k} (t :: k). Proxy t
Proxy :: Proxy flag)

class GRULayerStackRandomizable (flag :: Bool) inputSize hiddenSize numLayers directionality dtype device where
  gruLayerStackSample ::
    Proxy flag ->
    GRULayerStackSpec inputSize hiddenSize numLayers directionality dtype device ->
    IO (GRULayerStack inputSize hiddenSize numLayers directionality dtype device)

instance
  ( A.Randomizable
      (GRULayerSpec inputSize hiddenSize directionality dtype device)
      (GRULayer inputSize hiddenSize directionality dtype device)
  ) =>
  GRULayerStackRandomizable 'False inputSize hiddenSize 1 directionality dtype device
  where
  gruLayerStackSample :: Proxy 'False
-> GRULayerStackSpec
     inputSize hiddenSize 1 directionality dtype device
-> IO
     (GRULayerStack inputSize hiddenSize 1 directionality dtype device)
gruLayerStackSample Proxy 'False
_ GRULayerStackSpec
  inputSize hiddenSize 1 directionality dtype device
_ = forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayer inputSize hiddenSize directionality dtype device
-> GRULayerStack inputSize hiddenSize 1 directionality dtype device
GRULayer1 forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall spec f. Randomizable spec f => spec -> IO f
sample forall a b. (a -> b) -> a -> b
$ forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
GRULayerSpec @inputSize @hiddenSize @directionality @dtype @device)

instance
  ( 1 <= numLayers,
    A.Randomizable
      (GRULayerSpec (hiddenSize * NumberOfDirections directionality) hiddenSize directionality dtype device)
      (GRULayer (hiddenSize * NumberOfDirections directionality) hiddenSize directionality dtype device),
    A.Randomizable
      (GRULayerStackSpec inputSize hiddenSize (numLayers - 1) directionality dtype device)
      (GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device)
  ) =>
  GRULayerStackRandomizable 'True inputSize hiddenSize numLayers directionality dtype device
  where
  gruLayerStackSample :: Proxy 'True
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> IO
     (GRULayerStack
        inputSize hiddenSize numLayers directionality dtype device)
gruLayerStackSample Proxy 'True
_ GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
_ =
    forall (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)) (inputSize :: Natural)
       (numLayers :: Natural).
GRULayer
  (hiddenSize * NumberOfDirections directionality)
  hiddenSize
  directionality
  dtype
  device
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStack
     inputSize hiddenSize (numLayers + 1) directionality dtype device
GRULayerK
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall spec f. Randomizable spec f => spec -> IO f
sample forall a b. (a -> b) -> a -> b
$ forall (inputSize :: Natural) (hiddenSize :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerSpec inputSize hiddenSize directionality dtype device
GRULayerSpec @(hiddenSize * NumberOfDirections directionality) @hiddenSize @directionality @dtype @device)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> ( forall spec f. Randomizable spec f => spec -> IO f
sample
              @(GRULayerStackSpec inputSize hiddenSize (numLayers - 1) directionality dtype device)
              @(GRULayerStack inputSize hiddenSize (numLayers - 1) directionality dtype device)
              forall a b. (a -> b) -> a -> b
$ forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
GRULayerStackSpec
          )

instance
  ( 1 <= numLayers,
    (2 <=? numLayers) ~ flag,
    RandDTypeIsValid device dtype,
    KnownDType dtype,
    KnownDevice device,
    GRULayerStackRandomizable flag inputSize hiddenSize numLayers directionality dtype device
  ) =>
  Randomizable
    (GRULayerStackSpec inputSize hiddenSize numLayers directionality dtype device)
    (GRULayerStack inputSize hiddenSize numLayers directionality dtype device)
  where
  sample :: GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
-> IO
     (GRULayerStack
        inputSize hiddenSize numLayers directionality dtype device)
sample = forall (flag :: Bool) (inputSize :: Natural)
       (hiddenSize :: Natural) (numLayers :: Natural)
       (directionality :: RNNDirectionality) (dtype :: DType)
       (device :: (DeviceType, Natural)).
GRULayerStackRandomizable
  flag inputSize hiddenSize numLayers directionality dtype device =>
Proxy flag
-> GRULayerStackSpec
     inputSize hiddenSize numLayers directionality dtype device
-> IO
     (GRULayerStack
        inputSize hiddenSize numLayers directionality dtype device)
gruLayerStackSample (forall {k} (t :: k). Proxy t
Proxy :: Proxy flag)

newtype
  GRUSpec
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (numLayers :: Nat)
    (directionality :: RNNDirectionality)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  = GRUSpec DropoutSpec
  deriving (Int
-> GRUSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
Int
-> GRUSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
[GRUSpec
   inputSize hiddenSize numLayers directionality dtype device]
-> ShowS
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUSpec inputSize hiddenSize numLayers directionality dtype device
-> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GRUSpec
   inputSize hiddenSize numLayers directionality dtype device]
-> ShowS
$cshowList :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
[GRUSpec
   inputSize hiddenSize numLayers directionality dtype device]
-> ShowS
show :: GRUSpec inputSize hiddenSize numLayers directionality dtype device
-> String
$cshow :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUSpec inputSize hiddenSize numLayers directionality dtype device
-> String
showsPrec :: Int
-> GRUSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
$cshowsPrec :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
Int
-> GRUSpec
     inputSize hiddenSize numLayers directionality dtype device
-> ShowS
Show, forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)) x.
Rep
  (GRUSpec
     inputSize hiddenSize numLayers directionality dtype device)
  x
-> GRUSpec
     inputSize hiddenSize numLayers directionality dtype device
forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)) x.
GRUSpec inputSize hiddenSize numLayers directionality dtype device
-> Rep
     (GRUSpec
        inputSize hiddenSize numLayers directionality dtype device)
     x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)) x.
Rep
  (GRUSpec
     inputSize hiddenSize numLayers directionality dtype device)
  x
-> GRUSpec
     inputSize hiddenSize numLayers directionality dtype device
$cfrom :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)) x.
GRUSpec inputSize hiddenSize numLayers directionality dtype device
-> Rep
     (GRUSpec
        inputSize hiddenSize numLayers directionality dtype device)
     x
Generic)

data
  GRU
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (numLayers :: Nat)
    (directionality :: RNNDirectionality)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  where
  GRU ::
    (1 <= numLayers) =>
    { forall (numLayers :: Natural) (inputSize :: Natural)
       (hiddenSize :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
gru_layer_stack :: GRULayerStack inputSize hiddenSize numLayers directionality dtype device,
      forall (numLayers :: Natural) (inputSize :: Natural)
       (hiddenSize :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Dropout
gru_dropout :: Dropout
    } ->
    GRU inputSize hiddenSize numLayers directionality dtype device

deriving instance Show (GRU inputSize hiddenSize numLayers directionality dtype device)

instance
  (1 <= numLayers) =>
  Generic (GRU inputSize hiddenSize numLayers directionality dtype device)
  where
  type
    Rep (GRU inputSize hiddenSize numLayers directionality dtype device) =
      Rec0 (GRULayerStack inputSize hiddenSize numLayers directionality dtype device)
        :*: Rec0 Dropout
  from :: forall x.
GRU inputSize hiddenSize numLayers directionality dtype device
-> Rep
     (GRU inputSize hiddenSize numLayers directionality dtype device) x
from (GRU {Dropout
GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gru_dropout :: Dropout
gru_layer_stack :: GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gru_dropout :: forall (numLayers :: Natural) (inputSize :: Natural)
       (hiddenSize :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Dropout
gru_layer_stack :: forall (numLayers :: Natural) (inputSize :: Natural)
       (hiddenSize :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> GRULayerStack
     inputSize hiddenSize numLayers directionality dtype device
..}) = forall k i c (p :: k). c -> K1 i c p
K1 GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
gru_layer_stack forall k (f :: k -> Type) (g :: k -> Type) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall k i c (p :: k). c -> K1 i c p
K1 Dropout
gru_dropout
  to :: forall x.
Rep
  (GRU inputSize hiddenSize numLayers directionality dtype device) x
-> GRU inputSize hiddenSize numLayers directionality dtype device
to (K1 GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
layerStack :*: K1 Dropout
dropout) = forall (numLayers :: Natural) (inputSize :: Natural)
       (hiddenSize :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
(1 <= numLayers) =>
GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
-> Dropout
-> GRU inputSize hiddenSize numLayers directionality dtype device
GRU GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
layerStack Dropout
dropout

instance
  ( 1 <= numLayers,
    Parameterized (GRULayerStack inputSize hiddenSize numLayers directionality dtype device),
    HAppendFD
      (Parameters (GRULayerStack inputSize hiddenSize numLayers directionality dtype device))
      (Parameters Dropout)
      ( Parameters (GRULayerStack inputSize hiddenSize numLayers directionality dtype device)
          ++ Parameters Dropout
      )
  ) =>
  Parameterized (GRU inputSize hiddenSize numLayers directionality dtype device)

-- TODO: when we have cannonical initializers do this correctly:
-- https://github.com/pytorch/pytorch/issues/9221
-- https://discuss.pytorch.org/t/initializing-rnn-gru-and-gru-correctly/23605

-- | Helper to do xavier uniform initializations on weight matrices and
-- orthagonal initializations for the gates. (When implemented.)
xavierUniformGRU ::
  forall device dtype hiddenSize featureSize.
  ( KnownDType dtype,
    KnownNat hiddenSize,
    KnownNat featureSize,
    KnownDevice device,
    RandDTypeIsValid device dtype
  ) =>
  IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU :: forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (hiddenSize :: Natural) (featureSize :: Natural).
(KnownDType dtype, KnownNat hiddenSize, KnownNat featureSize,
 KnownDevice device, RandDTypeIsValid device dtype) =>
IO (Tensor device dtype '[3 * hiddenSize, featureSize])
xavierUniformGRU = do
  Tensor device dtype '[3 * hiddenSize, featureSize]
init <- forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(TensorOptions shape dtype device,
 RandDTypeIsValid device dtype) =>
IO (Tensor device dtype shape)
randn :: IO (Tensor device dtype '[3 * hiddenSize, featureSize])
  forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (shape :: [Natural]).
Tensor -> Tensor device dtype shape
UnsafeMkTensor
    forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Tensor -> Float -> [Int] -> IO Tensor
xavierUniformFIXME
      (forall t. Unnamed t => t -> Tensor
toDynamic Tensor device dtype '[3 * hiddenSize, featureSize]
init)
      (Float
5.0 forall a. Fractional a => a -> a -> a
/ Float
3)
      (forall (device :: (DeviceType, Natural)) (dtype :: DType)
       (shape :: [Natural]) t.
(TensorOptions shape dtype device,
 IsUnnamed t device dtype shape) =>
t -> [Int]
shape @device @dtype @'[3 * hiddenSize, featureSize] Tensor device dtype '[3 * hiddenSize, featureSize]
init)

instance
  ( KnownDType dtype,
    KnownDevice device,
    KnownNat inputSize,
    KnownNat hiddenSize,
    KnownNat (NumberOfDirections directionality),
    RandDTypeIsValid device dtype,
    A.Randomizable
      (GRULayerStackSpec inputSize hiddenSize numLayers directionality dtype device)
      (GRULayerStack inputSize hiddenSize numLayers directionality dtype device),
    1 <= numLayers
  ) =>
  A.Randomizable
    (GRUSpec inputSize hiddenSize numLayers directionality dtype device)
    (GRU inputSize hiddenSize numLayers directionality dtype device)
  where
  sample :: GRUSpec inputSize hiddenSize numLayers directionality dtype device
-> IO
     (GRU inputSize hiddenSize numLayers directionality dtype device)
sample (GRUSpec DropoutSpec
dropoutSpec) =
    forall (numLayers :: Natural) (inputSize :: Natural)
       (hiddenSize :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
(1 <= numLayers) =>
GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
-> Dropout
-> GRU inputSize hiddenSize numLayers directionality dtype device
GRU
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall spec f. Randomizable spec f => spec -> IO f
A.sample (forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRULayerStackSpec
  inputSize hiddenSize numLayers directionality dtype device
GRULayerStackSpec @inputSize @hiddenSize @numLayers @directionality @dtype @device)
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall spec f. Randomizable spec f => spec -> IO f
A.sample DropoutSpec
dropoutSpec

-- | A specification for a long, short-term memory layer.
data
  GRUWithInitSpec
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (numLayers :: Nat)
    (directionality :: RNNDirectionality)
    (initialization :: RNNInitialization)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  where
  -- | Weights drawn from Xavier-Uniform
  --   with zeros-value initialized biases and cell states.
  GRUWithZerosInitSpec ::
    forall inputSize hiddenSize numLayers directionality dtype device.
    GRUSpec inputSize hiddenSize numLayers directionality dtype device ->
    GRUWithInitSpec inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device
  -- | Weights drawn from Xavier-Uniform
  --   with zeros-value initialized biases
  --   and user-provided cell states.
  GRUWithConstInitSpec ::
    forall inputSize hiddenSize numLayers directionality dtype device.
    GRUSpec inputSize hiddenSize numLayers directionality dtype device ->
    -- | The initial values of the hidden state
    Tensor device dtype '[numLayers * NumberOfDirections directionality, hiddenSize] ->
    GRUWithInitSpec inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device
  -- | Weights drawn from Xavier-Uniform
  --   with zeros-value initialized biases
  --   and learned cell states.
  GRUWithLearnedInitSpec ::
    forall inputSize hiddenSize numLayers directionality dtype device.
    GRUSpec inputSize hiddenSize numLayers directionality dtype device ->
    -- | The initial (learnable)
    -- values of the hidden state
    Tensor device dtype '[numLayers * NumberOfDirections directionality, hiddenSize] ->
    GRUWithInitSpec inputSize hiddenSize numLayers directionality 'LearnedInitialization dtype device

deriving instance Show (GRUWithInitSpec inputSize hiddenSize numLayers directionality initialization dtype device)

-- | A long, short-term memory layer with either fixed initial
-- states for the memory cells and hidden state or learnable
-- inital states for the memory cells and hidden state.
data
  GRUWithInit
    (inputSize :: Nat)
    (hiddenSize :: Nat)
    (numLayers :: Nat)
    (directionality :: RNNDirectionality)
    (initialization :: RNNInitialization)
    (dtype :: D.DType)
    (device :: (D.DeviceType, Nat))
  where
  GRUWithConstInit ::
    forall inputSize hiddenSize numLayers directionality dtype device.
    { forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'ConstantInitialization
  dtype
  device
-> GRU inputSize hiddenSize numLayers directionality dtype device
gruWithConstInit_gru :: GRU inputSize hiddenSize numLayers directionality dtype device,
      forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'ConstantInitialization
  dtype
  device
-> Tensor
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithConstInit_h :: Tensor device dtype '[numLayers * NumberOfDirections directionality, hiddenSize]
    } ->
    GRUWithInit
      inputSize
      hiddenSize
      numLayers
      directionality
      'ConstantInitialization
      dtype
      device
  GRUWithLearnedInit ::
    forall inputSize hiddenSize numLayers directionality dtype device.
    { forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
-> GRU inputSize hiddenSize numLayers directionality dtype device
gruWithLearnedInit_gru :: GRU inputSize hiddenSize numLayers directionality dtype device,
      forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
-> Parameter
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithLearnedInit_h :: Parameter device dtype '[numLayers * NumberOfDirections directionality, hiddenSize]
    } ->
    GRUWithInit
      inputSize
      hiddenSize
      numLayers
      directionality
      'LearnedInitialization
      dtype
      device

deriving instance Show (GRUWithInit inputSize hiddenSize numLayers directionality initialization dtype device)

instance Generic (GRUWithInit inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device) where
  type
    Rep (GRUWithInit inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device) =
      Rec0 (GRU inputSize hiddenSize numLayers directionality dtype device)
        :*: Rec0 (Tensor device dtype '[numLayers * NumberOfDirections directionality, hiddenSize])
  from :: forall x.
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'ConstantInitialization
  dtype
  device
-> Rep
     (GRUWithInit
        inputSize
        hiddenSize
        numLayers
        directionality
        'ConstantInitialization
        dtype
        device)
     x
from (GRUWithConstInit {Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
GRU inputSize hiddenSize numLayers directionality dtype device
gruWithConstInit_h :: Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithConstInit_gru :: GRU inputSize hiddenSize numLayers directionality dtype device
gruWithConstInit_h :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'ConstantInitialization
  dtype
  device
-> Tensor
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithConstInit_gru :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'ConstantInitialization
  dtype
  device
-> GRU inputSize hiddenSize numLayers directionality dtype device
..}) = forall k i c (p :: k). c -> K1 i c p
K1 GRU inputSize hiddenSize numLayers directionality dtype device
gruWithConstInit_gru forall k (f :: k -> Type) (g :: k -> Type) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall k i c (p :: k). c -> K1 i c p
K1 Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithConstInit_h
  to :: forall x.
Rep
  (GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'ConstantInitialization
     dtype
     device)
  x
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'ConstantInitialization
     dtype
     device
to (K1 GRU inputSize hiddenSize numLayers directionality dtype device
gru :*: K1 Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h) = forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Tensor
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'ConstantInitialization
     dtype
     device
GRUWithConstInit GRU inputSize hiddenSize numLayers directionality dtype device
gru Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h

instance Generic (GRUWithInit inputSize hiddenSize numLayers directionality 'LearnedInitialization dtype device) where
  type
    Rep (GRUWithInit inputSize hiddenSize numLayers directionality 'LearnedInitialization dtype device) =
      Rec0 (GRU inputSize hiddenSize numLayers directionality dtype device)
        :*: Rec0 (Parameter device dtype '[numLayers * NumberOfDirections directionality, hiddenSize])
  from :: forall x.
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
-> Rep
     (GRUWithInit
        inputSize
        hiddenSize
        numLayers
        directionality
        'LearnedInitialization
        dtype
        device)
     x
from (GRUWithLearnedInit {Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
GRU inputSize hiddenSize numLayers directionality dtype device
gruWithLearnedInit_h :: Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithLearnedInit_gru :: GRU inputSize hiddenSize numLayers directionality dtype device
gruWithLearnedInit_h :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
-> Parameter
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithLearnedInit_gru :: forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
-> GRU inputSize hiddenSize numLayers directionality dtype device
..}) = forall k i c (p :: k). c -> K1 i c p
K1 GRU inputSize hiddenSize numLayers directionality dtype device
gruWithLearnedInit_gru forall k (f :: k -> Type) (g :: k -> Type) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall k i c (p :: k). c -> K1 i c p
K1 Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
gruWithLearnedInit_h
  to :: forall x.
Rep
  (GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'LearnedInitialization
     dtype
     device)
  x
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'LearnedInitialization
     dtype
     device
to (K1 GRU inputSize hiddenSize numLayers directionality dtype device
gru :*: K1 Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h) = forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Parameter
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'LearnedInitialization
     dtype
     device
GRUWithLearnedInit GRU inputSize hiddenSize numLayers directionality dtype device
gru Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h

instance
  ( Parameterized (GRU inputSize hiddenSize numLayers directionality dtype device),
    HAppendFD
      (Parameters (GRU inputSize hiddenSize numLayers directionality dtype device))
      '[]
      ( Parameters (GRU inputSize hiddenSize numLayers directionality dtype device) ++ '[]
      )
  ) =>
  Parameterized (GRUWithInit inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device)

instance
  ( Parameterized (GRU inputSize hiddenSize numLayers directionality dtype device),
    HAppendFD
      (Parameters (GRU inputSize hiddenSize numLayers directionality dtype device))
      '[Parameter device dtype '[numLayers * NumberOfDirections directionality, hiddenSize]]
      ( Parameters (GRU inputSize hiddenSize numLayers directionality dtype device)
          ++ '[Parameter device dtype '[numLayers * NumberOfDirections directionality, hiddenSize]]
      )
  ) =>
  Parameterized (GRUWithInit inputSize hiddenSize numLayers directionality 'LearnedInitialization dtype device)

instance
  ( KnownNat hiddenSize,
    KnownNat numLayers,
    KnownNat (NumberOfDirections directionality),
    KnownDType dtype,
    KnownDevice device,
    A.Randomizable
      (GRUSpec inputSize hiddenSize numLayers directionality dtype device)
      (GRU inputSize hiddenSize numLayers directionality dtype device)
  ) =>
  A.Randomizable
    (GRUWithInitSpec inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device)
    (GRUWithInit inputSize hiddenSize numLayers directionality 'ConstantInitialization dtype device)
  where
  sample :: GRUWithInitSpec
  inputSize
  hiddenSize
  numLayers
  directionality
  'ConstantInitialization
  dtype
  device
-> IO
     (GRUWithInit
        inputSize
        hiddenSize
        numLayers
        directionality
        'ConstantInitialization
        dtype
        device)
sample (GRUWithZerosInitSpec GRUSpec inputSize hiddenSize numLayers directionality dtype device
gruSpec) =
    forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Tensor
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'ConstantInitialization
     dtype
     device
GRUWithConstInit
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall spec f. Randomizable spec f => spec -> IO f
A.sample GRUSpec inputSize hiddenSize numLayers directionality dtype device
gruSpec
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
TensorOptions shape dtype device =>
Tensor device dtype shape
zeros
  sample (GRUWithConstInitSpec GRUSpec inputSize hiddenSize numLayers directionality dtype device
gruSpec Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h) =
    forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Tensor
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'ConstantInitialization
     dtype
     device
GRUWithConstInit
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall spec f. Randomizable spec f => spec -> IO f
A.sample GRUSpec inputSize hiddenSize numLayers directionality dtype device
gruSpec
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h

instance
  ( KnownNat hiddenSize,
    KnownNat numLayers,
    KnownNat (NumberOfDirections directionality),
    KnownDType dtype,
    KnownDevice device,
    A.Randomizable
      (GRUSpec inputSize hiddenSize numLayers directionality dtype device)
      (GRU inputSize hiddenSize numLayers directionality dtype device)
  ) =>
  A.Randomizable
    (GRUWithInitSpec inputSize hiddenSize numLayers directionality 'LearnedInitialization dtype device)
    (GRUWithInit inputSize hiddenSize numLayers directionality 'LearnedInitialization dtype device)
  where
  sample :: GRUWithInitSpec
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
-> IO
     (GRUWithInit
        inputSize
        hiddenSize
        numLayers
        directionality
        'LearnedInitialization
        dtype
        device)
sample s :: GRUWithInitSpec
  inputSize
  hiddenSize
  numLayers
  directionality
  'LearnedInitialization
  dtype
  device
s@(GRUWithLearnedInitSpec GRUSpec inputSize hiddenSize numLayers directionality dtype device
gruSpec Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h) =
    forall (inputSize :: Natural) (hiddenSize :: Natural)
       (numLayers :: Natural) (directionality :: RNNDirectionality)
       (dtype :: DType) (device :: (DeviceType, Natural)).
GRU inputSize hiddenSize numLayers directionality dtype device
-> Parameter
     device
     dtype
     '[numLayers * NumberOfDirections directionality, hiddenSize]
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     'LearnedInitialization
     dtype
     device
GRUWithLearnedInit
      forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> forall spec f. Randomizable spec f => spec -> IO f
A.sample GRUSpec inputSize hiddenSize numLayers directionality dtype device
gruSpec
      forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> (forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Tensor device dtype shape -> IO (Parameter device dtype shape)
makeIndependent forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (f :: Type -> Type) a. Applicative f => a -> f a
pure Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
h)

gruForward ::
  forall
    shapeOrder
    batchSize
    seqLen
    directionality
    initialization
    numLayers
    inputSize
    outputSize
    hiddenSize
    inputShape
    outputShape
    hcShape
    parameters
    tensorParameters
    dtype
    device.
  ( KnownNat (NumberOfDirections directionality),
    KnownNat numLayers,
    KnownNat batchSize,
    KnownNat hiddenSize,
    KnownRNNShapeOrder shapeOrder,
    KnownRNNDirectionality directionality,
    outputSize ~ (hiddenSize * NumberOfDirections directionality),
    inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
    outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
    hcShape ~ '[numLayers * NumberOfDirections directionality, batchSize, hiddenSize],
    parameters ~ Parameters (GRU inputSize hiddenSize numLayers directionality dtype device),
    Parameterized (GRU inputSize hiddenSize numLayers directionality dtype device),
    tensorParameters ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
    ATen.Castable (HList tensorParameters) [D.ATenTensor],
    HMap' ToDependent parameters tensorParameters
  ) =>
  Bool ->
  GRUWithInit
    inputSize
    hiddenSize
    numLayers
    directionality
    initialization
    dtype
    device ->
  Tensor device dtype inputShape ->
  ( Tensor device dtype outputShape,
    Tensor device dtype hcShape
  )
gruForward :: forall (shapeOrder :: RNNShapeOrder) (batchSize :: Natural)
       (seqLen :: Natural) (directionality :: RNNDirectionality)
       (initialization :: RNNInitialization) (numLayers :: Natural)
       (inputSize :: Natural) (outputSize :: Natural)
       (hiddenSize :: Natural) (inputShape :: [Natural])
       (outputShape :: [Natural]) (hcShape :: [Natural])
       (parameters :: [Type]) (tensorParameters :: [Type])
       (dtype :: DType) (device :: (DeviceType, Natural)).
(KnownNat (NumberOfDirections directionality), KnownNat numLayers,
 KnownNat batchSize, KnownNat hiddenSize,
 KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 parameters
 ~ Parameters
     (GRU inputSize hiddenSize numLayers directionality dtype device),
 Parameterized
   (GRU inputSize hiddenSize numLayers directionality dtype device),
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor],
 HMap' ToDependent parameters tensorParameters) =>
Bool
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     initialization
     dtype
     device
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gruForward Bool
dropoutOn (GRUWithConstInit gruModel :: GRU inputSize hiddenSize numLayers directionality dtype device
gruModel@(GRU GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
_ (Dropout Double
dropoutProb)) Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
hc) Tensor device dtype inputShape
input =
  forall {k} (shapeOrder :: RNNShapeOrder)
       (directionality :: RNNDirectionality) (numLayers :: Natural)
       (seqLen :: Natural) (batchSize :: Natural) (inputSize :: Natural)
       (outputSize :: Natural) (hiddenSize :: Natural)
       (inputShape :: [Natural]) (outputShape :: [Natural])
       (hcShape :: [Natural]) (tensorParameters :: [k]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(KnownNat numLayers, KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor]) =>
HList tensorParameters
-> Double
-> Bool
-> Tensor device dtype hcShape
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gru
    @shapeOrder
    @directionality
    @numLayers
    @seqLen
    @batchSize
    @inputSize
    @outputSize
    @hiddenSize
    @inputShape
    @outputShape
    @hcShape
    @tensorParameters
    @dtype
    @device
    (forall k f (xs :: [k]) (ys :: [k]).
HMap' f xs ys =>
f -> HList xs -> HList ys
hmap' ToDependent
ToDependent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f. Parameterized f => f -> HList (Parameters f)
flattenParameters forall a b. (a -> b) -> a -> b
$ GRU inputSize hiddenSize numLayers directionality dtype device
gruModel)
    Double
dropoutProb
    Bool
dropoutOn
    Tensor device dtype hcShape
hc'
    Tensor device dtype inputShape
input
  where
    hc' :: Tensor device dtype hcShape
hc' =
      forall (shape' :: [Natural]) (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(KnownShape shape', Numel shape ~ Numel shape') =>
Tensor device dtype shape -> Tensor device dtype shape'
reshape @hcShape
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (shape' :: [Natural]) (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(KnownShape shape', shape' ~ Broadcast shape shape') =>
Bool -> Tensor device dtype shape -> Tensor device dtype shape'
expand
          @'[batchSize, numLayers * NumberOfDirections directionality, hiddenSize]
          Bool
False -- TODO: What does the bool do?
        forall a b. (a -> b) -> a -> b
$ Tensor
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
hc
gruForward Bool
dropoutOn (GRUWithLearnedInit gruModel :: GRU inputSize hiddenSize numLayers directionality dtype device
gruModel@(GRU GRULayerStack
  inputSize hiddenSize numLayers directionality dtype device
_ (Dropout Double
dropoutProb)) Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
hc) Tensor device dtype inputShape
input =
  forall {k} (shapeOrder :: RNNShapeOrder)
       (directionality :: RNNDirectionality) (numLayers :: Natural)
       (seqLen :: Natural) (batchSize :: Natural) (inputSize :: Natural)
       (outputSize :: Natural) (hiddenSize :: Natural)
       (inputShape :: [Natural]) (outputShape :: [Natural])
       (hcShape :: [Natural]) (tensorParameters :: [k]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(KnownNat numLayers, KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor]) =>
HList tensorParameters
-> Double
-> Bool
-> Tensor device dtype hcShape
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gru
    @shapeOrder
    @directionality
    @numLayers
    @seqLen
    @batchSize
    @inputSize
    @outputSize
    @hiddenSize
    @inputShape
    @outputShape
    @hcShape
    @tensorParameters
    @dtype
    @device
    (forall k f (xs :: [k]) (ys :: [k]).
HMap' f xs ys =>
f -> HList xs -> HList ys
hmap' ToDependent
ToDependent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall f. Parameterized f => f -> HList (Parameters f)
flattenParameters forall a b. (a -> b) -> a -> b
$ GRU inputSize hiddenSize numLayers directionality dtype device
gruModel)
    Double
dropoutProb
    Bool
dropoutOn
    Tensor device dtype hcShape
hc'
    Tensor device dtype inputShape
input
  where
    hc' :: Tensor device dtype hcShape
hc' =
      forall (shape' :: [Natural]) (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(KnownShape shape', Numel shape ~ Numel shape') =>
Tensor device dtype shape -> Tensor device dtype shape'
reshape @hcShape
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (shape' :: [Natural]) (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
(KnownShape shape', shape' ~ Broadcast shape shape') =>
Bool -> Tensor device dtype shape -> Tensor device dtype shape'
expand
          @'[batchSize, numLayers * NumberOfDirections directionality, hiddenSize]
          Bool
False -- TODO: What does the bool do?
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (shape :: [Natural]) (dtype :: DType)
       (device :: (DeviceType, Natural)).
Parameter device dtype shape -> Tensor device dtype shape
toDependent
        forall a b. (a -> b) -> a -> b
$ Parameter
  device
  dtype
  '[numLayers * NumberOfDirections directionality, hiddenSize]
hc

gruForwardWithDropout,
  gruForwardWithoutDropout ::
    forall
      shapeOrder
      batchSize
      seqLen
      directionality
      initialization
      numLayers
      inputSize
      outputSize
      hiddenSize
      inputShape
      outputShape
      hcShape
      parameters
      tensorParameters
      dtype
      device.
    ( KnownNat (NumberOfDirections directionality),
      KnownNat numLayers,
      KnownNat batchSize,
      KnownNat hiddenSize,
      KnownRNNShapeOrder shapeOrder,
      KnownRNNDirectionality directionality,
      outputSize ~ (hiddenSize * NumberOfDirections directionality),
      inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
      outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
      hcShape ~ '[numLayers * NumberOfDirections directionality, batchSize, hiddenSize],
      parameters ~ Parameters (GRU inputSize hiddenSize numLayers directionality dtype device),
      Parameterized (GRU inputSize hiddenSize numLayers directionality dtype device),
      tensorParameters ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
      ATen.Castable (HList tensorParameters) [D.ATenTensor],
      HMap' ToDependent parameters tensorParameters
    ) =>
    GRUWithInit
      inputSize
      hiddenSize
      numLayers
      directionality
      initialization
      dtype
      device ->
    Tensor device dtype inputShape ->
    ( Tensor device dtype outputShape,
      Tensor device dtype hcShape
    )
-- ^ Forward propagate the `GRU` module and apply dropout on the outputs of each layer.
--
-- >>> input :: CPUTensor 'D.Float '[5,16,10] <- randn
-- >>> spec = GRUWithZerosInitSpec @10 @30 @3 @'Bidirectional @'D.Float @'( 'D.CPU, 0) (GRUSpec (DropoutSpec 0.5))
-- >>> model <- A.sample spec
-- >>> :t gruForwardWithDropout @'BatchFirst model input
-- gruForwardWithDropout @'BatchFirst model input
--   :: (Tensor '( 'D.CPU, 0) 'D.Float '[5, 16, 60],
--       Tensor '( 'D.CPU, 0) 'D.Float '[6, 5, 30])
-- >>> (a,b) = gruForwardWithDropout @'BatchFirst model input
-- >>> ((dtype a, shape a), (dtype b, shape b))
-- ((Float,[5,16,60]),(Float,[6,5,30]))
gruForwardWithDropout :: forall (shapeOrder :: RNNShapeOrder) (batchSize :: Natural)
       (seqLen :: Natural) (directionality :: RNNDirectionality)
       (initialization :: RNNInitialization) (numLayers :: Natural)
       (inputSize :: Natural) (outputSize :: Natural)
       (hiddenSize :: Natural) (inputShape :: [Natural])
       (outputShape :: [Natural]) (hcShape :: [Natural])
       (parameters :: [Type]) (tensorParameters :: [Type])
       (dtype :: DType) (device :: (DeviceType, Natural)).
(KnownNat (NumberOfDirections directionality), KnownNat numLayers,
 KnownNat batchSize, KnownNat hiddenSize,
 KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 parameters
 ~ Parameters
     (GRU inputSize hiddenSize numLayers directionality dtype device),
 Parameterized
   (GRU inputSize hiddenSize numLayers directionality dtype device),
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor],
 HMap' ToDependent parameters tensorParameters) =>
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  initialization
  dtype
  device
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gruForwardWithDropout =
  forall (shapeOrder :: RNNShapeOrder) (batchSize :: Natural)
       (seqLen :: Natural) (directionality :: RNNDirectionality)
       (initialization :: RNNInitialization) (numLayers :: Natural)
       (inputSize :: Natural) (outputSize :: Natural)
       (hiddenSize :: Natural) (inputShape :: [Natural])
       (outputShape :: [Natural]) (hcShape :: [Natural])
       (parameters :: [Type]) (tensorParameters :: [Type])
       (dtype :: DType) (device :: (DeviceType, Natural)).
(KnownNat (NumberOfDirections directionality), KnownNat numLayers,
 KnownNat batchSize, KnownNat hiddenSize,
 KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 parameters
 ~ Parameters
     (GRU inputSize hiddenSize numLayers directionality dtype device),
 Parameterized
   (GRU inputSize hiddenSize numLayers directionality dtype device),
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor],
 HMap' ToDependent parameters tensorParameters) =>
Bool
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     initialization
     dtype
     device
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gruForward
    @shapeOrder
    @batchSize
    @seqLen
    @directionality
    @initialization
    @numLayers
    @inputSize
    @outputSize
    @hiddenSize
    @inputShape
    @outputShape
    @hcShape
    @parameters
    @tensorParameters
    @dtype
    @device
    Bool
True
-- ^ Forward propagate the `GRU` module (without applying dropout on the outputs of each layer).
--
-- >>> input :: CPUTensor 'D.Float '[5,16,10] <- randn
-- >>> spec = GRUWithZerosInitSpec @10 @30 @3 @'Bidirectional @'D.Float @'( 'D.CPU, 0) (GRUSpec (DropoutSpec 0.5))
-- >>> model <- A.sample spec
-- >>> :t gruForwardWithoutDropout @'BatchFirst model input
-- gruForwardWithoutDropout @'BatchFirst model input
--   :: (Tensor '( 'D.CPU, 0) 'D.Float '[5, 16, 60],
--       Tensor '( 'D.CPU, 0) 'D.Float '[6, 5, 30])
-- >>> (a,b) = gruForwardWithoutDropout @'BatchFirst model input
-- >>> ((dtype a, shape a), (dtype b, shape b))
-- ((Float,[5,16,60]),(Float,[6,5,30]))
gruForwardWithoutDropout :: forall (shapeOrder :: RNNShapeOrder) (batchSize :: Natural)
       (seqLen :: Natural) (directionality :: RNNDirectionality)
       (initialization :: RNNInitialization) (numLayers :: Natural)
       (inputSize :: Natural) (outputSize :: Natural)
       (hiddenSize :: Natural) (inputShape :: [Natural])
       (outputShape :: [Natural]) (hcShape :: [Natural])
       (parameters :: [Type]) (tensorParameters :: [Type])
       (dtype :: DType) (device :: (DeviceType, Natural)).
(KnownNat (NumberOfDirections directionality), KnownNat numLayers,
 KnownNat batchSize, KnownNat hiddenSize,
 KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 parameters
 ~ Parameters
     (GRU inputSize hiddenSize numLayers directionality dtype device),
 Parameterized
   (GRU inputSize hiddenSize numLayers directionality dtype device),
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor],
 HMap' ToDependent parameters tensorParameters) =>
GRUWithInit
  inputSize
  hiddenSize
  numLayers
  directionality
  initialization
  dtype
  device
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gruForwardWithoutDropout =
  forall (shapeOrder :: RNNShapeOrder) (batchSize :: Natural)
       (seqLen :: Natural) (directionality :: RNNDirectionality)
       (initialization :: RNNInitialization) (numLayers :: Natural)
       (inputSize :: Natural) (outputSize :: Natural)
       (hiddenSize :: Natural) (inputShape :: [Natural])
       (outputShape :: [Natural]) (hcShape :: [Natural])
       (parameters :: [Type]) (tensorParameters :: [Type])
       (dtype :: DType) (device :: (DeviceType, Natural)).
(KnownNat (NumberOfDirections directionality), KnownNat numLayers,
 KnownNat batchSize, KnownNat hiddenSize,
 KnownRNNShapeOrder shapeOrder,
 KnownRNNDirectionality directionality,
 outputSize ~ (hiddenSize * NumberOfDirections directionality),
 inputShape ~ RNNShape shapeOrder seqLen batchSize inputSize,
 outputShape ~ RNNShape shapeOrder seqLen batchSize outputSize,
 hcShape
 ~ '[numLayers * NumberOfDirections directionality, batchSize,
     hiddenSize],
 parameters
 ~ Parameters
     (GRU inputSize hiddenSize numLayers directionality dtype device),
 Parameterized
   (GRU inputSize hiddenSize numLayers directionality dtype device),
 tensorParameters
 ~ GRUR inputSize hiddenSize numLayers directionality dtype device,
 Castable (HList tensorParameters) [ATenTensor],
 HMap' ToDependent parameters tensorParameters) =>
Bool
-> GRUWithInit
     inputSize
     hiddenSize
     numLayers
     directionality
     initialization
     dtype
     device
-> Tensor device dtype inputShape
-> (Tensor device dtype outputShape, Tensor device dtype hcShape)
gruForward
    @shapeOrder
    @batchSize
    @seqLen
    @directionality
    @initialization
    @numLayers
    @inputSize
    @outputSize
    @hiddenSize
    @inputShape
    @outputShape
    @hcShape
    @parameters
    @tensorParameters
    @dtype
    @device
    Bool
False