tensai

package module
v0.0.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 20, 2026 License: MIT Imports: 10 Imported by: 0

README

tensai

tensai - a tiny machine-learning framework in Go

tensai is a small machine-learning framework for learning and experiments. It implements forward passes, backpropagation, and optimization in pure Go with no external dependencies.

Features

  • Matrix operations - Matrix plus basic operations such as Dot, Add, T, and AddBias. Tensors are float32 (tensai.Float)
  • SIMD acceleration - AVX2 kernels written with Go's experimental simd/archsimd package: still pure Go, no cgo, no assembly files. Matmul, ReLU/LeakyReLU, Sigmoid/Tanh/Softmax (via a vectorized polynomial exp), GELU (via a vectorized erf), LayerNorm, and the Adam update are all 8-lane vectorized. Build with GOEXPERIMENT=simd on amd64 (Go 1.26 and 1.27 APIs both supported via build tags); every other build uses the portable fallbacks automatically
  • Low-allocation training - layers reuse their forward/backward scratch buffers across training steps (a full MLP step runs in ~29 allocations), so GC stays out of the training loop; Predict always returns freshly allocated results
  • Layers - Embedding, Dense, Conv2D, MaxPool2D, BatchNorm, LayerNorm, Dropout, plus ReLU, LeakyReLU, GELU, Sigmoid, Tanh, and Softmax activations
  • Loss functions - MeanSquaredError for regression, SoftmaxCrossEntropy for multi-class classification, and BinaryCrossEntropy for binary targets
  • Optimizers - momentum SGD, Adam, and AdamW (decoupled weight decay)
  • k-NN baseline - a KNN classifier whose distance matrix runs on the same SIMD matmul kernel; useful as a no-training baseline next to the networks
  • Sequential models - stack layers and run Compile -> Fit / FitStep -> Predict
  • Automatic differentiation - a micrograd-style reverse-mode autograd engine over matrices (Param / Input / Backward), for models that don't fit the Sequential mold; ToDot renders the computation graph for Graphviz
  • Recurrence and attention - RNNCell, LSTMCell, and single-head SelfAttention built on the autograd engine, with backpropagation through time handled automatically
  • Serialization - Save/Load (and SaveFile/LoadFile) round-trip trained Sequential parameters as JSON, including BatchNorm running statistics; SaveParams/LoadParams do the same for autograd parameters (RNN/LSTM/attention cells)
  • TFLite export - the encoding/tflite package marshals Sequential models (FP32, NHWC) into .tflite flatbuffers that run on the TFLite/LiteRT runtimes and go-tflite, with the FlatBuffers writer implemented in-tree — still no dependencies

Layout

go.mod              Module definition (github.com/mattn/tensai)
tensor.go           Matrix and vector operations
dot_simd.go         AVX2 matmul kernel (GOEXPERIMENT=simd, amd64)
dot_generic.go      Portable matmul kernel (all other builds)
kernels.go          Scalar bodies of the element-wise kernels
mathvec_simd.go     AVX2 element-wise kernels incl. vectorized exp
mathvec_generic.go  Portable element-wise kernel dispatchers
layer.go            Layer interface plus Dense and activations
conv.go             Conv2D and MaxPool2D layers
batchnorm.go        BatchNorm layer
dropout.go          Dropout layer
loss.go             Loss functions (MSE, SoftmaxCrossEntropy, BCE)
optimizer.go        Optimizers (SGD, Adam, AdamW)
model.go            Sequential model and training loop
autograd.go         Reverse-mode automatic differentiation (Node graph)
rnn.go              RNNCell / LSTMCell / SelfAttention on the autograd engine
serialize.go        Model parameter save/load (JSON)
tensai_test.go      Unit tests plus XOR convergence test
features_test.go    Gradient checks and tests for the newer layers
_example/helloworld Smallest possible program: add two values on the graph
_example/xor        Runnable XOR training example
_example/fizzbuzz   Runnable FizzBuzz classification example
_example/spiral     Runnable 3-class spiral classification example
_example/iris       Runnable Iris classification example
_example/mnist      Runnable MNIST classifier (-model dense, cnn, or knn) with save/load
_example/charrnn    Character-level LSTM text generation on the autograd engine
_example/plasma     Demoscene-style terminal plasma rendered by a neural network
_example/dot        Graphviz DOT export of the z = x + y graph

Usage

Regression: learn XOR with MSE
model := tensai.NewSequential()
model.Add(tensai.NewDense(8))
model.Add(&tensai.Tanh{})
model.Add(tensai.NewDense(1))
model.Add(&tensai.Sigmoid{})

model.Compile(2, tensai.MeanSquaredError{}, tensai.NewAdam(0.05))
model.Fit(inputs, targets, 5000)

pred, _ := model.Predict(inputs)
Classification: softmax + cross-entropy
model := tensai.NewSequential()
model.Add(tensai.NewDense(8))
model.Add(&tensai.ReLU{})
model.Add(tensai.NewDense(2)) // output width = number of classes

model.Compile(2, tensai.SoftmaxCrossEntropy{}, tensai.NewAdam(0.05))

SoftmaxCrossEntropy expects targets as an Mx1 matrix of class indices. Softmax is applied inside the loss, so Predict returns raw logits. Use argmax for classification.

Convolution, regularization, and saving
model := tensai.NewSequential()
model.Add(tensai.NewConv2D(28, 28, 1, 8, 3, 1, 1)) // inH, inW, inC, outC, kernel, stride, pad
model.Add(&tensai.ReLU{})
model.Add(tensai.NewMaxPool2D(28, 28, 8, 2))
model.Add(tensai.NewDense(64))
model.Add(tensai.NewBatchNorm())
model.Add(tensai.NewLeakyReLU(0.01))
model.Add(tensai.NewDropout(0.3))
model.Add(tensai.NewDense(10))

model.Compile(28*28, tensai.SoftmaxCrossEntropy{}, tensai.NewAdamW(0.001, 0.01))
model.Fit(inputs, targets, 10)

model.SaveFile("model.json")
// Later: build + Compile the same architecture, then
model.LoadFile("model.json")

Conv2D and MaxPool2D treat each row as a channel-major image: index = (channel*height + y)*width + x. Dropout and BatchNorm switch automatically between training behavior (inside Fit/FitStep) and inference behavior (inside Predict).

Embedding keeps the current matrix-only API: each input row is a token-id sequence, and the layer concatenates the looked-up embedding vectors across columns. For example, Compile(4, ...) plus NewEmbedding(vocab, 8) turns an Mx4 token-id matrix into an Mx32 dense feature matrix that can feed LayerNorm, GELU, and Dense.

TFLite export
import tensaitflite "github.com/mattn/tensai/encoding/tflite"

// after training:
err := tensaitflite.MarshalFile("model.tflite", model)

Supported layers: Dense, Conv2D (VALID/SAME padding), MaxPool2D, BatchNorm (folded into Mul+Add), Dropout (dropped), Softmax, and the ReLU/LeakyReLU/Sigmoid/Tanh activations. Exported convolutions follow TFLite's NHWC layout — feed the exported model NHWC input; weight reordering is handled during export. Outputs have been verified to match Predict to ~1e-7 relative error on the LiteRT interpreter (see encoding/tflite/verify_litert.py). Alias the import when combining with go-tflite, which also names its package tflite.

Automatic differentiation

When a model doesn't fit the Sequential mold (weight sharing, custom losses, exotic architectures), build the computation directly and let reverse-mode autodiff derive the gradients:

w1 := tensai.Param(tensai.RandomMatrix(2, 8, rng))
b1 := tensai.Param(tensai.NewMatrix(1, 8))
w2 := tensai.Param(tensai.RandomMatrix(8, 1, rng))
trainer := tensai.NewTrainer(tensai.NewAdam(0.05), w1, b1, w2)

for step := 0; step < 2000; step++ {
	loss := tensai.Input(x).MatMul(w1).AddRow(b1).Tanh().MatMul(w2).Sigmoid().MSELoss(y)
	trainer.Step(loss) // backward + update + zero grads, returns the loss value
}

For manual control, the pieces are still public: loss.Backward(), p.Grad, and tensai.ZeroGrads(params...).

The graph a loss node holds can be visualized: loss.ToDot() returns Graphviz DOT (label leaves with .Named("w1")), so go run ./_example/dot | dot -Tsvg > graph.svg draws the network the same way Gorgonia's encoding/dot does.

Graphs are built dynamically per step (define-by-run) and are single-use. Available ops: MatMul, Add, Sub, MulElem, Scale, AddRow, T, Softmax, ReLU, Sigmoid, Tanh, Sum, Mean, MSELoss, and SoftmaxCELoss. Shape mismatches panic during graph construction. Every op's gradient is verified against finite differences in the test suite.

Recurrent networks and attention

RNNCell, LSTMCell, and SelfAttention are built on the autograd engine, so unrolling a sequence is a plain Go loop and backpropagation through time comes for free:

cell := tensai.NewLSTMCell(inSize, hidden, rng)
wOut := tensai.Param(tensai.RandomMatrix(hidden, numClasses, rng))
bOut := tensai.Param(tensai.NewMatrix(1, numClasses))
trainer := tensai.NewTrainer(tensai.NewAdam(0.01), append(cell.Params(), wOut, bOut)...)

for step := 0; step < epochs; step++ {
	h, c := cell.InitState(batch)
	for _, x := range steps { // one (batch x inSize) matrix per time step
		h, c = cell.Step(tensai.Input(x), h, c)
	}
	logits := h.MatMul(wOut).AddRow(bOut)
	trainer.Step(logits.SoftmaxCELoss(labels))
}

SelfAttention operates on one (seqLen x inSize) sequence node: attn.Forward(x) computes softmax(Q*K^T/sqrt(d))*V with learned projections; the raw tensai.Attention(q, k, v) form is also exposed.

Autograd parameters are saved and restored positionally with tensai.SaveParamsFile("cell.json", cell.Params()...) / tensai.LoadParamsFile("cell.json", cell.Params()...) — build the same cell, then load.

Run

go run ./_example/helloworld
go run ./_example/xor
go run ./_example/fizzbuzz
go run ./_example/spiral
go run ./_example/iris
go run ./_example/charrnn
go run ./_example/plasma
go test ./...

# With the AVX2 SIMD kernel (Go 1.26+ / 1.27, amd64):
GOEXPERIMENT=simd go test ./...
GOEXPERIMENT=simd go test -bench=Dot .

The MNIST example downloads the standard IDX gzip files into _example/mnist/data when they are missing. Set MNIST_DIR to use another cache directory, and pass -model cnn for the convolutional variant (Conv2D/MaxPool2D/Dropout + AdamW); both trained variants finish by saving the model and re-scoring it after a reload. -model knn runs the no-training k-NN baseline instead — on the 5000-sample subset it scores ~91% against ~92% for the MLP and ~95% for the CNN:

go run ./_example/mnist
go run ./_example/mnist -model cnn
go run ./_example/mnist -model knn
MNIST_DIR=/path/to/mnist go run ./_example/mnist

The charrnn example trains a character-level LSTM on an embedded public-domain text, saves the parameters with SaveParamsFile, restores them into a fresh model, and generates a sample from the reloaded parameters.

The plasma example animates a demoscene-style plasma in the terminal where the plasma function is a randomly weighted network (a CPPN) evaluated for every pixel of every frame as one batch. The status line shows the per-frame network time, which makes it a live SIMD benchmark: 120x90 pixels runs at ~32 fps on the portable build and ~100 fps with GOEXPERIMENT=simd on the same machine. Try different -seed values for different effects.

Both raw IDX files and .gz variants are accepted.

SIMD Coverage

Where the AVX2 kernels apply today, and where they still could:

  • Matmul (Dot/DotInto) — used by Dense, Conv2D (im2col product), KNN distances, and autograd MatMul
  • ReLU / LeakyReLU forward & backward
  • Sigmoid / Tanh forward & backward (vectorized polynomial exp)
  • GELU forward & backward (vectorized erf)
  • LayerNorm forward & backward (vector row reductions)
  • Softmax / SoftmaxCrossEntropy exponentials and scaling
  • Adam / AdamW parameter update
  • Slice add & scale primitives (bias add, Embedding gradient scatter-add)
  • Transpose-free gradient matmul (DotTAInto) — Dense/Conv2D weight gradients no longer materialize input^T / im2col^T
  • Remaining transposes (T/TInto, now only weight matrices and autograd) — needs 8x8 block-and-shuffle
  • Softmax backward row dot products (layer and autograd)
  • MSE / BinaryCrossEntropy losses (BCE needs a vectorized log)
  • Autograd element-wise backward passes (gradients accumulate with +=, so they need dedicated fused kernels)
  • BatchNorm statistics (column-strided access needs a restructure)
  • MaxPool2D window scan
  • im2col / col2im gather-scatter (contiguous runs could use bulk copies)
  • SGD update

The unchecked items are ordered roughly by expected impact; none of them show up prominently in training profiles today.

Design Notes

  • All operations are batched. Inputs are MxN matrices, where M is the batch size and N is the feature dimension.
  • Embedding inputs are also matrices: values must be exact integer token ids stored in Float, and the embedding vectors are flattened across the row.
  • The Layer interface standardizes Forward, Backward, Params, and Grads, which keeps new layers such as convolution or dropout straightforward to add.
  • Dense weights use Glorot/He-style initialization to keep early training stable.
  • SoftmaxCrossEntropy subtracts the row maximum before softmax for numerical stability.

License

MIT

Author

Yasuhiro Matsumoto (a.k.a. mattn)

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func DotInto 

func DotInto(out, a, b *Matrix) error

DotInto computes out = a * b into an existing matrix, overwriting it.

func DotTAInto 

func DotTAInto(out, a, b *Matrix) error

DotTAInto computes out = a^T * b into an existing matrix, overwriting it, without materializing the transpose: a is read row by row and scattered into out with the same vector kernel Dot uses. Shapes: a is RxI, b is RxJ, out is IxJ.

func LoadParams 

func LoadParams(r io.Reader, params ...*Node) error

LoadParams restores parameter values saved by SaveParams. The parameters must be passed in the same order and have the same shapes as when saved.

func LoadParamsFile 

func LoadParamsFile(path string, params ...*Node) error

LoadParamsFile restores autograd parameters from a file written by SaveParamsFile.

func SaveParams 

func SaveParams(w io.Writer, params ...*Node) error

SaveParams writes the values of autograd parameters as JSON, in the given order. Pass the same parameter list a Trainer uses, e.g. SaveParams(w, cell.Params()...).

func SaveParamsFile 

func SaveParamsFile(path string, params ...*Node) error

SaveParamsFile writes autograd parameters to a JSON file.

func TInto 

func TInto(dst, src *Matrix) error

TInto writes the transpose of src into dst.

func ZeroGrads 

func ZeroGrads(nodes ...*Node)

ZeroGrads clears the gradients of the given nodes. Call it between training steps: Backward accumulates.

Types

type Adam 

type Adam struct {
	LR          Float
	Beta1       Float
	Beta2       Float
	Eps         Float
	WeightDecay Float
	// contains filtered or unexported fields
}

Adam optimizer. With WeightDecay > 0 it becomes AdamW: decay is decoupled from the gradient update and applied to weights only (never to biases).

func NewAdam 

func NewAdam(lr Float) *Adam

NewAdam returns an Adam optimizer with standard defaults.

func NewAdamW 

func NewAdamW(lr, weightDecay Float) *Adam

NewAdamW returns an Adam optimizer with decoupled weight decay (AdamW).

func (*Adam) Name 

func (a *Adam) Name() string

Name returns "adam".

func (*Adam) NewLayer 

func (a *Adam) NewLayer() int

NewLayer registers a layer and returns its index.

func (*Adam) Step 

func (a *Adam) Step(idx int, weights, gradW *Matrix, bias, gradB []Float)

Step updates parameters with the Adam rule.

func (*Adam) String 

func (a *Adam) String() string

String returns a human-readable description of an optimizer's config.

type BatchNorm 

type BatchNorm struct {
	Momentum Float // running-stats decay, default 0.9
	Eps      Float // numerical stability, default 1e-5
	// contains filtered or unexported fields
}

BatchNorm normalizes each feature column over the batch, then applies a learned scale (gamma) and shift (beta). During training it normalizes with batch statistics and maintains running estimates; during inference it uses the running estimates.

gamma is exposed as the layer's weights (a 1xC matrix) and beta as its bias, so optimizers update them like any other parameters.

func NewBatchNorm 

func NewBatchNorm() *BatchNorm

NewBatchNorm returns a BatchNorm layer with standard defaults.

func (*BatchNorm) Backward 

func (b *BatchNorm) Backward(gradOutput *Matrix) (*Matrix, error)

func (*BatchNorm) Forward 

func (b *BatchNorm) Forward(input *Matrix) (*Matrix, error)

func (*BatchNorm) Grads 

func (b *BatchNorm) Grads() (*Matrix, []Float)

func (*BatchNorm) Init 

func (b *BatchNorm) Init(inputCols int, _ *rand.Rand) (int, error)

func (*BatchNorm) Params 

func (b *BatchNorm) Params() (*Matrix, []Float)

func (*BatchNorm) RunningStats 

func (b *BatchNorm) RunningStats() (mean, variance []Float)

RunningStats returns the running mean and variance estimates used at inference time, for exporters.

func (*BatchNorm) SetParams 

func (b *BatchNorm) SetParams(weights *Matrix, bias []Float) error

type BinaryCrossEntropy 

type BinaryCrossEntropy struct{}

BinaryCrossEntropy computes the average binary cross-entropy between predicted probabilities in (0,1) and 0/1 targets of the same shape. Pair it with a Sigmoid output layer.

func (BinaryCrossEntropy) Loss 

func (BinaryCrossEntropy) Loss(pred, target *Matrix) (Float, *Matrix, error)

Loss returns the average binary cross-entropy and its gradient.

func (BinaryCrossEntropy) LossInto 

func (BinaryCrossEntropy) LossInto(pred, target, grad *Matrix) (Float, error)

LossInto writes the BCE gradient into grad and returns the average loss.

func (BinaryCrossEntropy) Name 

func (BinaryCrossEntropy) Name() string

Name returns "bce".

type Conv2D 

type Conv2D struct {
	// contains filtered or unexported fields
}

Conv2D is a 2D convolution layer. Because the framework moves data as flat MxN matrices, each sample row must be laid out channel-major: index = (channel*height + y)*width + x. The output uses the same layout.

The convolution is computed as a matrix product over an im2col expansion, so it reuses the tuned Dot kernel. Weights are stored as an (inC*kernel*kernel) x outC matrix.

func NewConv2D 

func NewConv2D(inH, inW, inC, outC, kernel, stride, pad int) *Conv2D

NewConv2D returns a convolution layer for inH x inW inputs with inC channels, producing outC channels with a square kernel.

func (*Conv2D) Backward 

func (c *Conv2D) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Conv2D) Forward 

func (c *Conv2D) Forward(input *Matrix) (*Matrix, error)

func (*Conv2D) Grads 

func (c *Conv2D) Grads() (*Matrix, []Float)

func (*Conv2D) Init 

func (c *Conv2D) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Conv2D) Params 

func (c *Conv2D) Params() (*Matrix, []Float)

func (*Conv2D) SetParams 

func (c *Conv2D) SetParams(weights *Matrix, bias []Float) error

func (*Conv2D) Shape 

func (c *Conv2D) Shape() (inH, inW, inC, outC, kernel, stride, pad int)

Shape reports the layer's spatial configuration, for exporters.

type Dense 

type Dense struct {
	// contains filtered or unexported fields
}

Dense is a fully-connected layer: y = x*W + b.

func NewDense 

func NewDense(outCols int) *Dense

NewDense returns a Dense layer with the given output size.

func (*Dense) Backward 

func (d *Dense) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Dense) Forward 

func (d *Dense) Forward(input *Matrix) (*Matrix, error)

func (*Dense) Grads 

func (d *Dense) Grads() (*Matrix, []Float)

func (*Dense) Init 

func (d *Dense) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Dense) Params 

func (d *Dense) Params() (*Matrix, []Float)

func (*Dense) SetParams 

func (d *Dense) SetParams(weights *Matrix, bias []Float) error

type Dropout 

type Dropout struct {
	Rate Float
	// contains filtered or unexported fields
}

Dropout randomly zeroes elements during training with probability Rate and scales the survivors by 1/(1-Rate) ("inverted dropout"), so inference is a plain pass-through with no rescaling.

func NewDropout 

func NewDropout(rate Float) *Dropout

NewDropout returns a Dropout layer that drops the given fraction of activations during training. Rate must be in [0, 1).

func (*Dropout) Backward 

func (d *Dropout) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Dropout) Forward 

func (d *Dropout) Forward(input *Matrix) (*Matrix, error)

func (*Dropout) Grads 

func (d *Dropout) Grads() (*Matrix, []Float)

func (*Dropout) Init 

func (d *Dropout) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Dropout) Params 

func (d *Dropout) Params() (*Matrix, []Float)

func (*Dropout) SetParams 

func (d *Dropout) SetParams(*Matrix, []Float) error

type Embedding 

type Embedding struct {
	// contains filtered or unexported fields
}

Embedding looks up a learned vector for each token id in the input row and concatenates the vectors across columns.

func NewEmbedding 

func NewEmbedding(vocabSize, dim int) *Embedding

NewEmbedding returns a trainable embedding table of shape vocabSize x dim.

func (*Embedding) Backward 

func (e *Embedding) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Embedding) Forward 

func (e *Embedding) Forward(input *Matrix) (*Matrix, error)

func (*Embedding) Grads 

func (e *Embedding) Grads() (*Matrix, []Float)

func (*Embedding) Init 

func (e *Embedding) Init(inputCols int, rng *rand.Rand) (int, error)

func (*Embedding) Params 

func (e *Embedding) Params() (*Matrix, []Float)

func (*Embedding) SetParams 

func (e *Embedding) SetParams(weights *Matrix, bias []Float) error

type Float 

type Float = float32

Float is the element type of every tensor. float32 halves memory traffic versus float64 and enables the 8-lane AVX2 kernel (see dot_simd.go); its ~7 decimal digits are plenty for neural-network training.

type GELU 

type GELU struct {
	// contains filtered or unexported fields
}

GELU activation: f(x) = 0.5*x*(1+erf(x/sqrt(2))).

func (*GELU) Backward 

func (g *GELU) Backward(gradOutput *Matrix) (*Matrix, error)

func (*GELU) Forward 

func (g *GELU) Forward(input *Matrix) (*Matrix, error)

func (*GELU) Grads 

func (a *GELU) Grads() (*Matrix, []Float)

func (*GELU) Init 

func (g *GELU) Init(inputCols int, _ *rand.Rand) (int, error)

func (*GELU) Params 

func (a *GELU) Params() (*Matrix, []Float)

func (*GELU) SetParams 

func (a *GELU) SetParams(*Matrix, []Float) error

type KNN 

type KNN struct {
	K int
	// contains filtered or unexported fields
}

KNN is a k-nearest-neighbors classifier. It is a lazy learner: Fit just stores the training data, and Predict ranks neighbors by squared Euclidean distance. The distance computation is reduced to one matrix product per chunk (||a-b||^2 = ||a||^2 + ||b||^2 - 2*a.b), so it runs on the same tuned Dot kernel as the neural networks.

func NewKNN 

func NewKNN(k int) *KNN

NewKNN returns a classifier that votes among the k nearest neighbors.

func (*KNN) Fit 

func (k *KNN) Fit(inputs, targets *Matrix) error

Fit stores the training set. Targets must be an Mx1 matrix of class indices, as with SoftmaxCrossEntropy.

func (*KNN) Predict 

func (k *KNN) Predict(inputs *Matrix) (*Matrix, error)

Predict returns an (inputs.Rows x classes) matrix of neighbor vote counts; take the argmax per row for the predicted class. Row sums equal K.

type LSTMCell 

type LSTMCell struct {
	// One (Wx, Wh, B) triple per gate: forget, input, output, candidate.
	Wxf, Whf, Bf *Node
	Wxi, Whi, Bi *Node
	Wxo, Who, Bo *Node
	Wxg, Whg, Bg *Node
}

LSTMCell is a long short-term memory cell with forget/input/output gates.

func NewLSTMCell 

func NewLSTMCell(inSize, hidden int, rng *rand.Rand) *LSTMCell

NewLSTMCell returns a randomly initialized LSTM cell. Forget-gate biases start at 1 so early training defaults to remembering.

func (*LSTMCell) InitState 

func (c *LSTMCell) InitState(batch int) (*Node, *Node)

InitState returns zero (hidden, cell) states for the given batch size.

func (*LSTMCell) Params 

func (c *LSTMCell) Params() []*Node

Params returns the cell's trainable parameters, for NewTrainer.

func (*LSTMCell) Step 

func (c *LSTMCell) Step(x, h, cell *Node) (*Node, *Node)

Step consumes one time step with the previous (hidden, cell) state and returns the next (hidden, cell) state.

type Layer 

type Layer interface {
	// Init configures parameters using the given RNG and input width.
	Init(inputCols int, rng *rand.Rand) (outputCols int, err error)

	// Forward computes activations given the input batch.
	Forward(input *Matrix) (*Matrix, error)

	// Backward computes the gradient with respect to the layer input,
	// given the gradient with respect to the layer output.
	Backward(gradOutput *Matrix) (*Matrix, error)

	// Grads returns the parameter gradients accumulated during the last
	// backward pass, in the order [weights, bias]. Layers without
	// parameters return nil.
	Grads() (*Matrix, []Float)

	// Params returns the current parameters [weights, bias].
	Params() (*Matrix, []Float)

	// SetParams replaces the parameters [weights, bias].
	SetParams(weights *Matrix, bias []Float) error
}

Layer is a single differentiable stage of a Sequential model. Forward and Backward are batched: inputs/outputs are MxN matrices where M is the batch size and N is the feature dimension.

type LayerNorm 

type LayerNorm struct {
	// contains filtered or unexported fields
}

LayerNorm normalizes each row over its feature dimension and applies a learnable affine transform.

func NewLayerNorm 

func NewLayerNorm() *LayerNorm

NewLayerNorm returns a LayerNorm with the default epsilon.

func (*LayerNorm) Backward 

func (l *LayerNorm) Backward(gradOutput *Matrix) (*Matrix, error)

func (*LayerNorm) Forward 

func (l *LayerNorm) Forward(input *Matrix) (*Matrix, error)

func (*LayerNorm) Grads 

func (l *LayerNorm) Grads() (*Matrix, []Float)

func (*LayerNorm) Init 

func (l *LayerNorm) Init(inputCols int, _ *rand.Rand) (int, error)

func (*LayerNorm) Params 

func (l *LayerNorm) Params() (*Matrix, []Float)

func (*LayerNorm) SetParams 

func (l *LayerNorm) SetParams(weights *Matrix, bias []Float) error

type LeakyReLU 

type LeakyReLU struct {
	Alpha Float
	// contains filtered or unexported fields
}

LeakyReLU activation: f(x) = x for x > 0, alpha*x otherwise.

func NewLeakyReLU 

func NewLeakyReLU(alpha Float) *LeakyReLU

NewLeakyReLU returns a LeakyReLU with the given negative-side slope.

func (*LeakyReLU) Backward 

func (l *LeakyReLU) Backward(gradOutput *Matrix) (*Matrix, error)

func (*LeakyReLU) Forward 

func (l *LeakyReLU) Forward(input *Matrix) (*Matrix, error)

func (*LeakyReLU) Grads 

func (a *LeakyReLU) Grads() (*Matrix, []Float)

func (*LeakyReLU) Init 

func (l *LeakyReLU) Init(inputCols int, _ *rand.Rand) (int, error)

func (*LeakyReLU) Params 

func (a *LeakyReLU) Params() (*Matrix, []Float)

func (*LeakyReLU) SetParams 

func (a *LeakyReLU) SetParams(*Matrix, []Float) error

type Loss 

type Loss interface {
	// Loss returns the average loss and the per-element gradient dL/dpred.
	Loss(pred, target *Matrix) (Float, *Matrix, error)
	// Name is a short identifier for logging.
	Name() string
}

Loss is a differentiable loss function operating on a prediction batch and a target batch of the same shape. It returns the scalar loss and the gradient of the loss with respect to the predictions.

type Matrix 

type Matrix struct {
	Rows int
	Cols int
	Data []Float
}

Matrix is a row-major 2D tensor of Float.

func Add 

func Add(a, b *Matrix) (*Matrix, error)

Add returns a + b (element-wise). Shapes must match.

func AddBias 

func AddBias(a *Matrix, bias []Float) (*Matrix, error)

AddBias adds a 1xCols bias vector to every row of a.

func Dot 

func Dot(a, b *Matrix) (*Matrix, error)

Dot computes the matrix product a * b.

func NewMatrix 

func NewMatrix(rows, cols int) *Matrix

NewMatrix creates a matrix filled with zeros.

func NewMatrixFromSlice 

func NewMatrixFromSlice(rows, cols int, data []Float) (*Matrix, error)

NewMatrixFromSlice creates a rows x cols matrix from row-major data.

func RandomMatrix 

func RandomMatrix(rows, cols int, rng *rand.Rand) *Matrix

RandomMatrix fills a matrix with samples from a normal distribution scaled by the Glorot/Bengio gain for the given fan-in / fan-out.

func (*Matrix) At 

func (m *Matrix) At(r, c int) Float

At returns the element at (r, c).

func (*Matrix) Row 

func (m *Matrix) Row(r int) []Float

Row returns a copy of row r as a slice.

func (*Matrix) Scale 

func (m *Matrix) Scale(s Float)

Scale multiplies every element by s, in place.

func (*Matrix) Set 

func (m *Matrix) Set(r, c int, v Float)

Set sets the element at (r, c).

func (*Matrix) SetRow 

func (m *Matrix) SetRow(r int, vals []Float) error

SetRow copies vals into row r.

func (*Matrix) T 

func (m *Matrix) T() *Matrix

T returns the transpose of the matrix.

func (*Matrix) Validate 

func (m *Matrix) Validate() error

Validate returns an error if the matrix data length is inconsistent.

type MaxPool2D 

type MaxPool2D struct {
	// contains filtered or unexported fields
}

MaxPool2D downsamples each channel by taking the maximum over non-overlapping size x size windows. It expects the same channel-major layout as Conv2D.

func NewMaxPool2D 

func NewMaxPool2D(inH, inW, channels, size int) *MaxPool2D

NewMaxPool2D returns a max-pooling layer with stride equal to size.

func (*MaxPool2D) Backward 

func (p *MaxPool2D) Backward(gradOutput *Matrix) (*Matrix, error)

func (*MaxPool2D) Forward 

func (p *MaxPool2D) Forward(input *Matrix) (*Matrix, error)

func (*MaxPool2D) Grads 

func (p *MaxPool2D) Grads() (*Matrix, []Float)

func (*MaxPool2D) Init 

func (p *MaxPool2D) Init(inputCols int, _ *rand.Rand) (int, error)

func (*MaxPool2D) Params 

func (p *MaxPool2D) Params() (*Matrix, []Float)

func (*MaxPool2D) SetParams 

func (p *MaxPool2D) SetParams(*Matrix, []Float) error

func (*MaxPool2D) Shape 

func (p *MaxPool2D) Shape() (inH, inW, channels, size int)

Shape reports the layer's spatial configuration, for exporters.

type MeanSquaredError 

type MeanSquaredError struct{}

MeanSquaredError computes the average of squared differences.

func (MeanSquaredError) Loss 

func (MeanSquaredError) Loss(pred, target *Matrix) (Float, *Matrix, error)

Loss returns the mean squared error and its gradient.

func (MeanSquaredError) LossInto 

func (MeanSquaredError) LossInto(pred, target, grad *Matrix) (Float, error)

LossInto writes the MSE gradient into grad and returns the average loss.

func (MeanSquaredError) Name 

func (MeanSquaredError) Name() string

Name returns "mse".

type Model 

type Model interface {
	Predict(input *Matrix) (*Matrix, error)
}

Model is a trained, ready-to-predict network.

type Node 

type Node struct {
	Value *Matrix
	Grad  *Matrix
	// contains filtered or unexported fields
}

Node is a matrix-valued node in a dynamically built computation graph for reverse-mode automatic differentiation. Build the forward computation by chaining operations, then call Backward on the (scalar) result to fill Grad on every Param node that contributed to it.

Unlike the Layer API, shape mismatches panic: graph construction errors are programming errors, and error returns would make chaining unusable. 

w := tensai.Param(tensai.RandomMatrix(2, 8, rng))
b := tensai.Param(tensai.NewMatrix(1, 8))
loss := tensai.Input(x).MatMul(w).AddRow(b).ReLU().MSELoss(y)
loss.Backward()
// w.Grad and b.Grad now hold dLoss/dw and dLoss/db.

func Attention 

func Attention(q, k, v *Node) *Node

Attention computes scaled dot-product attention softmax(q*k^T/sqrt(d))*v for a single sequence, where q, k, v are (seqLen x d) nodes.

func Input 

func Input(m *Matrix) *Node

Input wraps a matrix as a constant graph leaf. No gradient is computed for it.

func Param 

func Param(m *Matrix) *Node

Param wraps a matrix as a trainable graph leaf. Backward accumulates its gradient into Grad.

func (*Node) Add 

func (n *Node) Add(o *Node) *Node

Add returns the element-wise sum n + o. Shapes must match.

func (*Node) AddRow 

func (n *Node) AddRow(row *Node) *Node

AddRow adds a 1xN row node to every row of n (bias broadcast).

func (*Node) Backward 

func (n *Node) Backward()

Backward runs reverse-mode differentiation from n, which should be a scalar (1x1) loss. Gradients accumulate into the Grad field of every contributing Param node.

func (*Node) MSELoss 

func (n *Node) MSELoss(target *Matrix) *Node

MSELoss returns the scalar mean squared error against a constant target.

func (*Node) MatMul 

func (n *Node) MatMul(o *Node) *Node

MatMul returns the matrix product n * o.

func (*Node) Mean 

func (n *Node) Mean() *Node

Mean reduces n to a 1x1 scalar averaging all elements.

func (*Node) MulElem 

func (n *Node) MulElem(o *Node) *Node

MulElem returns the element-wise (Hadamard) product. Shapes must match.

func (*Node) Named 

func (n *Node) Named(name string) *Node

Named sets a display name shown by ToDot and returns the node.

func (*Node) ReLU 

func (n *Node) ReLU() *Node

ReLU applies max(0, x) element-wise.

func (*Node) Scalar 

func (n *Node) Scalar() Float

Scalar returns the value of a 1x1 node.

func (*Node) Scale 

func (n *Node) Scale(s Float) *Node

Scale returns n multiplied by the scalar s.

func (*Node) Sigmoid 

func (n *Node) Sigmoid() *Node

Sigmoid applies 1/(1+e^-x) element-wise.

func (*Node) Softmax 

func (n *Node) Softmax() *Node

Softmax normalizes each row of n into a probability distribution.

func (*Node) SoftmaxCELoss 

func (n *Node) SoftmaxCELoss(target *Matrix) *Node

SoftmaxCELoss returns the scalar softmax cross-entropy against integer class labels (an Mx1 matrix of class indices), matching the SoftmaxCrossEntropy loss used by Sequential models.

func (*Node) Sub 

func (n *Node) Sub(o *Node) *Node

Sub returns the element-wise difference n - o. Shapes must match.

func (*Node) Sum 

func (n *Node) Sum() *Node

Sum reduces n to a 1x1 scalar by summing all elements.

func (*Node) T 

func (n *Node) T() *Node

T returns the transpose of n.

func (*Node) Tanh 

func (n *Node) Tanh() *Node

Tanh applies tanh(x) element-wise.

func (*Node) ToDot 

func (n *Node) ToDot() string

ToDot renders the computation graph rooted at n in Graphviz DOT format, in the spirit of Gorgonia's encoding/dot. Pipe it through the dot tool to get an image: 

go run ./_example/dot | dot -Tsvg > graph.svg

Leaves are drawn as boxes (Param blue, Input gray) and operations as rounded nodes; every node shows its shape. Use Named to label leaves.

type Optimizer 

type Optimizer interface {
	// Step applies one update to the given parameters using their gradients.
	// It is called once per parameterized layer.
	Step(idx int, weights, gradW *Matrix, bias, gradB []Float)
	// NewLayer registers a new parameterized layer and returns its index.
	NewLayer() int
	// Name returns a short identifier.
	Name() string
}

Optimizer updates a set of (weights, bias) parameter pairs using their gradients. One Optimizer instance is shared by the model; each parameterized layer gets its own state buffer inside the optimizer.

type RNNCell 

type RNNCell struct {
	Wx *Node // in x hidden
	Wh *Node // hidden x hidden
	B  *Node // 1 x hidden
}

RNNCell is a simple (Elman) recurrent cell: h' = tanh(x*Wx + h*Wh + b).

func NewRNNCell 

func NewRNNCell(inSize, hidden int, rng *rand.Rand) *RNNCell

NewRNNCell returns a randomly initialized RNN cell.

func (*RNNCell) InitState 

func (c *RNNCell) InitState(batch int) *Node

InitState returns a zero hidden state for the given batch size.

func (*RNNCell) Params 

func (c *RNNCell) Params() []*Node

Params returns the cell's trainable parameters, for NewTrainer.

func (*RNNCell) Step 

func (c *RNNCell) Step(x, h *Node) *Node

Step consumes one time step and the previous hidden state, returning the next hidden state.

type ReLU 

type ReLU struct {
	// contains filtered or unexported fields
}

ReLU activation: f(x) = max(0, x).

func (*ReLU) Backward 

func (r *ReLU) Backward(gradOutput *Matrix) (*Matrix, error)

func (*ReLU) Forward 

func (r *ReLU) Forward(input *Matrix) (*Matrix, error)

func (*ReLU) Grads 

func (a *ReLU) Grads() (*Matrix, []Float)

func (*ReLU) Init 

func (r *ReLU) Init(inputCols int, _ *rand.Rand) (int, error)

func (*ReLU) Params 

func (a *ReLU) Params() (*Matrix, []Float)

func (*ReLU) SetParams 

func (a *ReLU) SetParams(*Matrix, []Float) error

type SGD 

type SGD struct {
	LR       Float
	Momentum Float
	// contains filtered or unexported fields
}

SGD is stochastic gradient descent with optional momentum.

func NewSGD 

func NewSGD(lr, momentum Float) *SGD

NewSGD returns an SGD optimizer. Momentum of 0 disables momentum.

func (*SGD) Name 

func (s *SGD) Name() string

Name returns "sgd".

func (*SGD) NewLayer 

func (s *SGD) NewLayer() int

NewLayer registers a layer and returns its index.

func (*SGD) Step 

func (s *SGD) Step(idx int, weights, gradW *Matrix, bias, gradB []Float)

Step updates parameters with standard (momentum) SGD.

func (*SGD) String 

func (s *SGD) String() string

String returns a human-readable description of an optimizer's config.

type SelfAttention 

type SelfAttention struct {
	Wq, Wk, Wv *Node // inSize x dModel
}

SelfAttention is a single-head self-attention block with learned query, key, and value projections. It operates on one sequence at a time: the input is a (seqLen x inSize) node.

func NewSelfAttention 

func NewSelfAttention(inSize, dModel int, rng *rand.Rand) *SelfAttention

NewSelfAttention returns a randomly initialized self-attention block.

func (*SelfAttention) Forward 

func (a *SelfAttention) Forward(x *Node) *Node

Forward applies self-attention to a (seqLen x inSize) sequence, returning a (seqLen x dModel) sequence.

func (*SelfAttention) Params 

func (a *SelfAttention) Params() []*Node

Params returns the block's trainable parameters, for NewTrainer.

type Sequential 

type Sequential struct {
	// contains filtered or unexported fields
}

Sequential stacks layers and runs forward/backward passes.

func NewSequential 

func NewSequential() *Sequential

NewSequential returns an empty Sequential model. optimizer and loss are configured via Compile.

func (*Sequential) Add 

func (s *Sequential) Add(layer Layer) *Sequential

Add appends a layer to the network. Layers are added in forward order.

func (*Sequential) Compile 

func (s *Sequential) Compile(inputCols int, loss Loss, optimizer Optimizer) error

Compile wires the loss and optimizer and initializes all parameters. inputCols is the number of features in a single input row.

func (*Sequential) Fit 

func (s *Sequential) Fit(input, target *Matrix, epochs int) error

Fit trains the model for the given number of epochs over the dataset. If epochs > 1 the full dataset is reused each epoch (full-batch by default; callers can pass minibatches to FitStep directly for finer control).

func (*Sequential) FitStep 

func (s *Sequential) FitStep(input, target *Matrix) (Float, error)

FitStep performs one forward + loss + backward + update pass for a batch and returns the average loss for that batch.

func (*Sequential) Layers 

func (s *Sequential) Layers() []Layer

Layers returns the layers in forward order, for tools that walk the model structure (e.g. format exporters).

func (*Sequential) Load 

func (s *Sequential) Load(r io.Reader) error

Load restores parameters saved by Save into a model compiled with the same architecture.

func (*Sequential) LoadFile 

func (s *Sequential) LoadFile(path string) error

LoadFile restores parameters from a file written by SaveFile.

func (*Sequential) Predict 

func (s *Sequential) Predict(input *Matrix) (*Matrix, error)

Predict runs a forward pass with no gradient tracking.

func (*Sequential) Save 

func (s *Sequential) Save(w io.Writer) error

Save writes the model's parameters as JSON. The architecture itself is not stored: Load must be called on a model built and compiled with the same layers.

func (*Sequential) SaveFile 

func (s *Sequential) SaveFile(path string) error

SaveFile writes the model's parameters to a JSON file.

type Sigmoid 

type Sigmoid struct {
	// contains filtered or unexported fields
}

Sigmoid activation: f(x) = 1 / (1 + e^-x).

func (*Sigmoid) Backward 

func (s *Sigmoid) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Sigmoid) Forward 

func (s *Sigmoid) Forward(input *Matrix) (*Matrix, error)

func (*Sigmoid) Grads 

func (a *Sigmoid) Grads() (*Matrix, []Float)

func (*Sigmoid) Init 

func (s *Sigmoid) Init(inputCols int, _ *rand.Rand) (int, error)

func (*Sigmoid) Params 

func (a *Sigmoid) Params() (*Matrix, []Float)

func (*Sigmoid) SetParams 

func (a *Sigmoid) SetParams(*Matrix, []Float) error

type Softmax 

type Softmax struct {
	// contains filtered or unexported fields
}

Softmax normalizes each row into a probability distribution. Unlike the element-wise activations its backward pass couples all columns of a row. Note that SoftmaxCrossEntropy already applies softmax internally; use this layer only when the model output itself must be probabilities (e.g. with a custom loss).

func (*Softmax) Backward 

func (s *Softmax) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Softmax) Forward 

func (s *Softmax) Forward(input *Matrix) (*Matrix, error)

func (*Softmax) Grads 

func (a *Softmax) Grads() (*Matrix, []Float)

func (*Softmax) Init 

func (s *Softmax) Init(inputCols int, _ *rand.Rand) (int, error)

func (*Softmax) Params 

func (a *Softmax) Params() (*Matrix, []Float)

func (*Softmax) SetParams 

func (a *Softmax) SetParams(*Matrix, []Float) error

type SoftmaxCrossEntropy 

type SoftmaxCrossEntropy struct{}

SoftmaxCrossEntropy combines softmax + cross-entropy with integer class labels. Targets must be an Mx1 matrix whose entries are class indices.

func (SoftmaxCrossEntropy) Loss 

func (SoftmaxCrossEntropy) Loss(pred, target *Matrix) (Float, *Matrix, error)

Loss returns the average negative log-likelihood of the target classes and the combined softmax-cross-entropy gradient (pred - onehot) / batch.

func (SoftmaxCrossEntropy) LossInto 

func (SoftmaxCrossEntropy) LossInto(pred, target, grad *Matrix) (Float, error)

LossInto writes the softmax-cross-entropy gradient into grad.

func (SoftmaxCrossEntropy) Name 

func (SoftmaxCrossEntropy) Name() string

Name returns "softmax_ce".

type Tanh 

type Tanh struct {
	// contains filtered or unexported fields
}

Tanh activation: f(x) = tanh(x).

func (*Tanh) Backward 

func (t *Tanh) Backward(gradOutput *Matrix) (*Matrix, error)

func (*Tanh) Forward 

func (t *Tanh) Forward(input *Matrix) (*Matrix, error)

func (*Tanh) Grads 

func (a *Tanh) Grads() (*Matrix, []Float)

func (*Tanh) Init 

func (t *Tanh) Init(inputCols int, _ *rand.Rand) (int, error)

func (*Tanh) Params 

func (a *Tanh) Params() (*Matrix, []Float)

func (*Tanh) SetParams 

func (a *Tanh) SetParams(*Matrix, []Float) error

type Trainer 

type Trainer struct {
	// contains filtered or unexported fields
}

Trainer owns the optimizer bookkeeping for a set of autograd parameters, so a training step is just building the loss graph and calling Step. 

trainer := tensai.NewTrainer(tensai.NewAdam(0.05), w1, b1, w2, b2)
for step := 0; step < 2000; step++ {
	loss := forward(x).MSELoss(y)
	trainer.Step(loss)
}

func NewTrainer 

func NewTrainer(opt Optimizer, params ...*Node) *Trainer

NewTrainer registers the parameters with the optimizer and returns a Trainer that updates them.

func (*Trainer) Step 

func (t *Trainer) Step(loss *Node) Float

Step runs backward from the scalar loss, applies one optimizer update to every parameter, clears the gradients, and returns the loss value.

Source Files

View all Source files

Directories

Path Synopsis
_example
charrnn command
Command charrnn trains a character-level LSTM on a small embedded corpus and generates text from it.
 Command charrnn trains a character-level LSTM on a small embedded corpus and generates text from it.
dot command
Command dot prints the computation graph of z = x + y in Graphviz DOT format — tensai's equivalent of Gorgonia's encoding/dot example.
 Command dot prints the computation graph of z = x + y in Graphviz DOT format — tensai's equivalent of Gorgonia's encoding/dot example.
fizzbuzz command
helloworld command
Command helloworld is the smallest possible tensai program: build a computation graph that adds two values, evaluate it, and differentiate it — tensai's equivalent of Gorgonia's hello world.
 Command helloworld is the smallest possible tensai program: build a computation graph that adds two values, evaluate it, and differentiate it — tensai's equivalent of Gorgonia's hello world.
iris command
mnist command
plasma command
Command plasma renders a demoscene-style plasma effect in the terminal — except the plasma function is a neural network.
 Command plasma renders a demoscene-style plasma effect in the terminal — except the plasma function is a neural network.
spiral command
xor command
encoding
tflite
Package tflite marshals trained tensai Sequential models into the TensorFlow Lite FlatBuffers format (FP32, batch size 1), so they can run on the TFLite / LiteRT runtimes — including from Go via github.com/mattn/go-tflite (alias one of the packages when importing both, e.g.
 Package tflite marshals trained tensai Sequential models into the TensorFlow Lite FlatBuffers format (FP32, batch size 1), so they can run on the TFLite / LiteRT runtimes — including from Go via github.com/mattn/go-tflite (alias one of the packages when importing both, e.g.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.