sequence

Documentation

Overview

Package sequence provides sequence modeling layers such as State Space Models.

Index

• type S4
  • func NewS4[T tensor.Float](name string, engine compute.Engine[T], ops numeric.Arithmetic[T], ...) (*S4[T], error)
  • func (s *S4[T]) Attributes() map[string]interface{}
  • func (s *S4[T]) Backward(_ context.Context, _ types.BackwardMode, ...) ([]*tensor.TensorNumeric[T], error)
  • func (s *S4[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)
  • func (s *S4[T]) OpType() string
  • func (s *S4[T]) OutputShape() []int
  • func (s *S4[T]) Parameters() []*graph.Parameter[T]

Types

type S4

type S4[T tensor.Float] struct {
	// contains filtered or unexported fields
}

S4 implements a diagonal State Space Model (S4D variant).

The continuous-time state space model is:

x'(t) = A x(t) + B u(t)
y(t)  = C x(t) + D u(t)

With diagonal A, the discrete-time equations become element-wise:

x_k = a * x_{k-1} + b * u_k
y_k = sum(c * x_k) + d * u_k

where a = exp(dt * A_diag) ensures stability when A_diag < 0.

Input shape: [batch, seq_len, input_dim] Output shape: [batch, seq_len, input_dim]

Parameters (per input dimension, state_dim internal states):

A_log [input_dim, state_dim] - log(-A), parameterizing stable eigenvalues
B     [input_dim, state_dim] - input-to-state projection
C     [input_dim, state_dim] - state-to-output projection
D     [input_dim]            - skip connection

func NewS4

func NewS4[T tensor.Float](
	name string,
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	inputDim, stateDim int,
) (*S4[T], error)

NewS4 creates a new S4 layer with HiPPO-inspired initialization.

func (*S4[T]) Attributes

func (s *S4[T]) Attributes() map[string]interface{}

Attributes returns the layer attributes.

func (*S4[T]) Backward

func (s *S4[T]) Backward(_ context.Context, _ types.BackwardMode, outputGradient *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes gradients using full backpropagation through time (BPTT).

The adjoint equation for the hidden state is:

dL/dx_k = dL/dy_k * C + dL/dx_{k+1} * A_disc

We iterate backward from the last timestep, accumulating gradients for all parameters.

func (*S4[T]) Forward

func (s *S4[T]) Forward(ctx context.Context, inputs ...*tensor.TensorNumeric[T]) (*tensor.TensorNumeric[T], error)

Forward runs the diagonal SSM scan over the sequence.

Input: [batch, seq_len, input_dim] Output: [batch, seq_len, input_dim]

All arithmetic is routed through engine primitives so the computation graph is fully traceable by the tracing compiler.

func (*S4[T]) OpType

func (s *S4[T]) OpType() string

OpType returns the operation type.

func (*S4[T]) OutputShape

func (s *S4[T]) OutputShape() []int

OutputShape returns the output shape.

func (*S4[T]) Parameters

func (s *S4[T]) Parameters() []*graph.Parameter[T]

Parameters returns all trainable parameters.