JANETCell
LuxRecurrentLayers.JANETCell
— TypeJANETCell(in_dims => out_dims;
use_bias=true, train_state=false, train_memory=false,
init_bias=nothing, init_weight=nothing, init_recurrent_weight=nothing,
init_state=zeros32, init_memory=zeros32, beta=1.0)
Equations
\[\begin{aligned} \mathbf{s}(t) &= \mathbf{W}_{ih}^{f} \mathbf{x}(t) + \mathbf{b}_{ih}^{f} + \mathbf{W}_{hh}^{f} \mathbf{h}(t-1) + \mathbf{b}_{hh}^{f}, \\ \tilde{\mathbf{c}}(t) &= \tanh\left( \mathbf{W}_{ih}^{c} \mathbf{x}(t) + \mathbf{b}_{ih}^{c} + \mathbf{W}_{hh}^{c} \mathbf{h}(t-1) + \mathbf{b}_{hh}^{c} \right), \\ \mathbf{c}(t) &= \sigma(\mathbf{s}(t)) \circ \mathbf{c}(t-1) + \left(1 - \sigma(\mathbf{s}(t) - \beta)\right) \circ \tilde{\mathbf{c}}(t), \\ \mathbf{h}(t) &= \mathbf{c}(t) \end{aligned}\]
Arguments
in_dims
: Input Dimensionout_dims
: Output (Hidden State & Memory) Dimension
Keyword Arguments
use_bias
: Flag to use bias in the computation. Default set totrue
.train_state
: Flag to set the initial hidden state as trainable. Default set tofalse
.train_memory
: Flag to set the initial memory state as trainable. Default set tofalse
.init_bias
: Initializer for input-to-hidden biases $\mathbf{b}_{ih}^{f}$ and $\mathbf{b}_{ih}^{c}$. Must be a tuple of 2 functions, e.g.,(glorot_uniform, kaiming_uniform)
. If a single functionfn
is provided, it is expanded to(fn, fn)
. If set tonothing
, each bias is initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default isnothing
.init_recurrent_bias
: Initializer for hidden-to-hidden biases $\mathbf{b}_{hh}^{f}$ and $\mathbf{b}_{hh}^{c}$. Must be a tuple of 2 functions, e.g.,(glorot_uniform, kaiming_uniform)
. If a single functionfn
is provided, it is expanded to(fn, fn)
. If set tonothing
, each bias is initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default isnothing
.init_weight
: Initializer for input-to-hidden weights $\mathbf{W}_{ih}^{f}$ and $\mathbf{W}_{ih}^{c}$. Must be a tuple of 2 functions, e.g.,(glorot_uniform, kaiming_uniform)
. If a single functionfn
is provided, it is expanded to(fn, fn)
. If set tonothing
, each weight is initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default isnothing
.init_recurrent_weight
: Initializer for hidden-to-hidden weights $\mathbf{W}_{hh}^{f}$ and $\mathbf{W}_{hh}^{c}$. Must be a tuple of 2 functions, e.g.,(glorot_uniform, kaiming_uniform)
. If a single functionfn
is provided, it is expanded to(fn, fn)
. If set tonothing
, each weight is initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default isnothing
.init_state
: Initializer for hidden state. Default set tozeros32
.init_memory
: Initializer for memory. Default set tozeros32
.beta
: Control parameter over the input data flow. Default is1.0
.
Inputs
- Case 1a: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set tofalse
,train_memory
is set tofalse
- Creates a hidden state usinginit_state
, hidden memory usinginit_memory
and proceeds to Case 2. - Case 1b: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set totrue
,train_memory
is set tofalse
- Repeatshidden_state
vector from the parameters to match the shape ofx
, creates hidden memory usinginit_memory
and proceeds to Case 2. - Case 1c: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set tofalse
,train_memory
is set totrue
- Creates a hidden state usinginit_state
, repeats the memory vector from parameters to match the shape ofx
and proceeds to Case 2. - Case 1d: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set totrue
,train_memory
is set totrue
- Repeats the hidden state and memory vectors from the parameters to match the shape ofx
and proceeds to Case 2. - Case 2: Tuple
(x, (h, c))
is provided, then the output and a tuple containing the updated hidden state and memory is returned.
Returns
Tuple Containing
- Output $h_{new}$ of shape
(out_dims, batch_size)
- Tuple containing new hidden state $h_{new}$ and new memory $c_{new}$
- Output $h_{new}$ of shape
Updated model state
Parameters
weight_ih
: Concatenated weights mapping from input to hidden units $\{ \mathbf{W}_{ih}^{f}, \mathbf{W}_{ih}^{c} \}$ The functions provided ininit_weight
are applied in order: the first function initializes $\mathbf{W}_{ih}^{f}$, the second initializes $\mathbf{W}_{ih}^{c}$.weight_hh
: Concatenated weights mapping from hidden state to hidden units $\{ \mathbf{W}_{hh}^{f}, \mathbf{W}_{hh}^{c} \}$ The functions provided ininit_recurrent_weight
are applied in order: the first function initializes $\mathbf{W}_{hh}^{f}$, the second initializes $\mathbf{W}_{hh}^{c}$.bias_ih
: Concatenated input-to-hidden bias vectors (ifuse_bias=true
) $\{ \mathbf{b}_{ih}^{f}, \mathbf{b}_{ih}^{c} \}$ The functions provided ininit_bias
are applied in order: the first function initializes $\mathbf{b}_{ih}^{f}$, the second initializes $\mathbf{b}_{ih}^{c}$.bias_hh
: Concatenated hidden-to-hidden bias vectors (ifuse_bias=true
) $\{ \mathbf{b}_{hh}^{f}, \mathbf{b}_{hh}^{c} \}$ The functions provided ininit_recurrent_bias
are applied in order: the first function initializes $\mathbf{b}_{hh}^{f}$, the second initializes $\mathbf{b}_{hh}^{c}$.hidden_state
: Initial hidden state vector (not present iftrain_state=false
)memory
: Initial memory vector (not present iftrain_memory=false
)
States
rng
: Controls the randomness (if any) in the initial state generation