{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Torch.NN.Recurrent.Cell.GRU where

import GHC.Generics
import Torch

data GRUSpec = GRUSpec
  { GRUSpec -> Int
inputSize :: Int,
    GRUSpec -> Int
hiddenSize :: Int
  }
  deriving (GRUSpec -> GRUSpec -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GRUSpec -> GRUSpec -> Bool
$c/= :: GRUSpec -> GRUSpec -> Bool
== :: GRUSpec -> GRUSpec -> Bool
$c== :: GRUSpec -> GRUSpec -> Bool
Eq, Int -> GRUSpec -> ShowS
[GRUSpec] -> ShowS
GRUSpec -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GRUSpec] -> ShowS
$cshowList :: [GRUSpec] -> ShowS
show :: GRUSpec -> String
$cshow :: GRUSpec -> String
showsPrec :: Int -> GRUSpec -> ShowS
$cshowsPrec :: Int -> GRUSpec -> ShowS
Show)

data GRUCell = GRUCell
  { GRUCell -> Parameter
weightsIH :: Parameter,
    GRUCell -> Parameter
weightsHH :: Parameter,
    GRUCell -> Parameter
biasIH :: Parameter,
    GRUCell -> Parameter
biasHH :: Parameter
  }
  deriving (forall x. Rep GRUCell x -> GRUCell
forall x. GRUCell -> Rep GRUCell x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep GRUCell x -> GRUCell
$cfrom :: forall x. GRUCell -> Rep GRUCell x
Generic, Int -> GRUCell -> ShowS
[GRUCell] -> ShowS
GRUCell -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GRUCell] -> ShowS
$cshowList :: [GRUCell] -> ShowS
show :: GRUCell -> String
$cshow :: GRUCell -> String
showsPrec :: Int -> GRUCell -> ShowS
$cshowsPrec :: Int -> GRUCell -> ShowS
Show)

gruCellForward ::
  -- | cell parameters
  GRUCell ->
  -- | input
  Tensor ->
  -- | hidden
  Tensor ->
  -- | output
  Tensor
gruCellForward :: GRUCell -> Tensor -> Tensor -> Tensor
gruCellForward GRUCell {Parameter
biasHH :: Parameter
biasIH :: Parameter
weightsHH :: Parameter
weightsIH :: Parameter
biasHH :: GRUCell -> Parameter
biasIH :: GRUCell -> Parameter
weightsHH :: GRUCell -> Parameter
weightsIH :: GRUCell -> Parameter
..} Tensor
input Tensor
hidden =
  Tensor -> Tensor -> Tensor -> Tensor -> Tensor -> Tensor -> Tensor
gruCell Tensor
weightsIH' Tensor
weightsHH' Tensor
biasIH' Tensor
biasHH' Tensor
hidden Tensor
input
  where
    weightsIH' :: Tensor
weightsIH' = Parameter -> Tensor
toDependent Parameter
weightsIH
    weightsHH' :: Tensor
weightsHH' = Parameter -> Tensor
toDependent Parameter
weightsHH
    biasIH' :: Tensor
biasIH' = Parameter -> Tensor
toDependent Parameter
biasIH
    biasHH' :: Tensor
biasHH' = Parameter -> Tensor
toDependent Parameter
biasHH

instance Parameterized GRUCell

instance Randomizable GRUSpec GRUCell where
  sample :: GRUSpec -> IO GRUCell
sample GRUSpec {Int
hiddenSize :: Int
inputSize :: Int
hiddenSize :: GRUSpec -> Int
inputSize :: GRUSpec -> Int
..} = do
    -- https://pytorch.org/docs/stable/generated/torch.nn.GRUCell.html
    Parameter
weightsIH' <- Tensor -> IO Parameter
makeIndependent forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Tensor -> Tensor
initScale forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> IO Tensor
randIO' [Int
3 forall a. Num a => a -> a -> a
* Int
hiddenSize, Int
inputSize]
    Parameter
weightsHH' <- Tensor -> IO Parameter
makeIndependent forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Tensor -> Tensor
initScale forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> IO Tensor
randIO' [Int
3 forall a. Num a => a -> a -> a
* Int
hiddenSize, Int
hiddenSize]
    Parameter
biasIH' <- Tensor -> IO Parameter
makeIndependent forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Tensor -> Tensor
initScale forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> IO Tensor
randIO' [Int
3 forall a. Num a => a -> a -> a
* Int
hiddenSize]
    Parameter
biasHH' <- Tensor -> IO Parameter
makeIndependent forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Tensor -> Tensor
initScale forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> IO Tensor
randIO' [Int
3 forall a. Num a => a -> a -> a
* Int
hiddenSize]
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      GRUCell
        { weightsIH :: Parameter
weightsIH = Parameter
weightsIH',
          weightsHH :: Parameter
weightsHH = Parameter
weightsHH',
          biasIH :: Parameter
biasIH = Parameter
biasIH',
          biasHH :: Parameter
biasHH = Parameter
biasHH'
        }
    where
      scale :: Float
scale = forall a. Floating a => a -> a
Prelude.sqrt forall a b. (a -> b) -> a -> b
$ Float
1.0 forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
hiddenSize :: Float
      initScale :: Tensor -> Tensor
initScale = forall a. Scalar a => a -> Tensor -> Tensor
subScalar Float
scale forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Scalar a => a -> Tensor -> Tensor
mulScalar Float
scale forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Scalar a => a -> Tensor -> Tensor
mulScalar (Float
2.0 :: Float)