Struct GruCell 

pub struct GruCell { /* private fields */ }

One PyTorch-form GRU transition with explicit hidden state.

Gate order is PyTorch’s reset, update, new (r, z, n), and the reset gate multiplies the hidden-side candidate contribution after its own matrix product – r * (W_hn h + b_hn), PyTorch’s form – rather than the paper’s r * h before any projection: r = sigmoid(W_ir x + b_ir + W_hr h + b_hr), z = sigmoid(W_iz x + b_iz + W_hz h + b_hz), n = tanh(W_in x + b_in + r * (W_hn h + b_hn)), h' = (1 - z) * n + z * h. Mirrors LstmCell’s single combined bias (rather than PyTorch’s separate bias_ih/bias_hh) and its Xavier-uniform initialization (rather than PyTorch’s uniform(-1/sqrt(hidden), 1/sqrt(hidden))); Axis’s own Lstm does not match PyTorch’s parameter layout or init either, so GruCell follows LstmCell’s convention instead of introducing a second one. bias defaults to true (PyTorch’s default); .bias(false) matches bias=False.

Implementations

impl GruCell

pub fn new(input: Axis, hidden: Dim) -> Result<Self>

pub fn bias(self, bias: bool) -> Self

Toggle the combined bias parameter. Has no effect once build has already allocated parameters; call it before build. Defaults to true.

pub fn zero_state(&self, input: &Tensor) -> Result<Tensor>

pub fn step(&self, input: &Tensor, state: &Tensor) -> Result<Tensor>

Trait Implementations

impl Module for GruCell

fn output_shape(&self, input: &Shape) -> Result<Shape>

fn build(&mut self, input: &Shape, device: &Device, seed: u64) -> Result<Shape>

Infer extents, validate contracts, allocate and initialize parameters; never run forward. Rebuilding with compatible extents/device preserves parameters and ignores the new seed.

fn forward(&self, input: &Tensor) -> Result<Tensor>

Evaluation semantics: no mode flag on self, no randomness beyond what forward itself already defines. This never changes with the arrival of forward_training.

fn named_parameters(&self) -> Vec<(String, Parameter)>

Structural paths identify slots; ParamId identifies shared storage within this process.

fn forward_training( &self, input: &Tensor, pass: &mut TrainingPass, ) -> Result<Tensor>

Training-mode forward pass, driven by an explicit TrainingPass rather than a flag on self. The default implementation ignores pass and calls Module::forward, which is exactly today’s behavior for every module without a train/eval distinction. Override this (never forward, which must stay evaluation semantics) to add explicit randomness or a train/eval difference, such as Dropout. A container that holds child modules (Sequential, or any future one) overrides this to thread the same pass through every child in order, so their draws stay distinct within one step.

fn parameter(&self, name: &str) -> Result<Parameter>

fn parameters(&self) -> Vec<Parameter>

fn zero_grad(&self)

fn named_states(&self) -> Vec<(String, State)>

Persistent non-parameter tensors this module owns across training steps – BatchNorm-style running statistics. Same slot-path convention as Module::named_parameters: a container (Sequential or any future one) aggregates every child’s own states under "{index}.{name}". Defaults empty, exactly like named_parameters, since most modules have no persistent state. A State here is never written directly; only TrainingPass::commit writes one, after staging in forward_training.

fn state(&self, name: &str) -> Result<State>

fn states(&self) -> Vec<State>