normalization

package
v1.56.0 Latest Latest
Warning

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

 Go to latest Published: Jul 3, 2026 License: Apache-2.0 Imports: 12 Imported by: 0

Documentation

Overview

Package normalization provides normalization layers for neural networks.

Stability: stable

Package normalization provides various normalization layers for neural networks.

Package normalization provides various normalization layers for neural networks.

Package normalization provides normalization layers for the Zerfoo model.

Package normalization provides normalization layers for the Zerfoo model.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func BuildBatchNormalization added in v0.2.1 

func BuildBatchNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	_ string,
	_ map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildBatchNormalization constructs a BatchNormalization layer for the registry. Reads "epsilon" (float32 or float64) from attributes; defaults to 1e-5.

func BuildGroupNormalization added in v1.54.0 

func BuildGroupNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	_ string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildGroupNormalization constructs a GroupNormalization layer for the registry. Reads "num_groups" (int/int64, default 32) and "epsilon" (default 1e-5).

func BuildLayerNormalization added in v0.2.1 

func BuildLayerNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildLayerNormalization constructs a LayerNormalization node. It resolves scale (gamma) and bias (beta) parameters using several naming conventions derived from the node name, and reads epsilon from attributes.

func BuildRMSNorm added in v0.2.0 

func BuildRMSNorm[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildRMSNorm constructs an RMSNorm node from the provided parameter and epsilon attribute.

func BuildSimplifiedLayerNormalization added in v0.2.0 

func BuildSimplifiedLayerNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildSimplifiedLayerNormalization constructs a SimplifiedLayerNormalization node, attempting multiple common naming patterns to resolve the gain/weight parameter.

func BuildSkipSimplifiedLayerNormalization added in v0.2.0 

func BuildSkipSimplifiedLayerNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	name string,
	params map[string]*graph.Parameter[T],
	attributes map[string]interface{},
) (graph.Node[T], error)

BuildSkipSimplifiedLayerNormalization constructs a SkipSimplifiedLayerNormalization node, resolving the gain/weight parameter using several naming conventions.

Types

type BatchNormalization added in v0.2.1 

type BatchNormalization[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

BatchNormalization implements batch normalization with optional training support. Forward expects 5 inputs: X, scale, B, mean, var (in that order). Formula per element at channel c: 

y = scale[c] * (X - mean[c]) / sqrt(var[c] + epsilon) + B[c]

scale, B, mean, and var must have shape [C] matching X's channel dimension. All operations use Engine primitives so they appear in the ExecutionPlan instruction tape.

When scale and bias are set as graph.Parameter (via NewBatchNormalizationWithParams), Backward() computes and accumulates gradients for scale, bias, and input X.

func NewBatchNormalization added in v0.2.1 

func NewBatchNormalization[T tensor.Numeric](engine compute.Engine[T], ops numeric.Arithmetic[T], epsilon T) *BatchNormalization[T]

NewBatchNormalization creates an inference-only BatchNormalization layer.

func NewBatchNormalizationWithParams added in v1.16.0 

func NewBatchNormalizationWithParams[T tensor.Numeric](engine compute.Engine[T], ops numeric.Arithmetic[T], epsilon T, scale, bias *graph.Parameter[T]) *BatchNormalization[T]

NewBatchNormalizationWithParams creates a BatchNormalization layer with learnable scale and bias parameters for training.

func (*BatchNormalization[T]) Attributes added in v0.2.1 

func (b *BatchNormalization[T]) Attributes() map[string]interface{}

Attributes returns the layer configuration.

func (*BatchNormalization[T]) Backward added in v0.2.1 

func (b *BatchNormalization[T]) Backward(ctx context.Context, _ types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes gradients for scale, bias, and input X.

For inference-mode BatchNorm (pre-computed mean/var): 

dBias  = sum(dOut, dims except channel)
dScale = sum(dOut * normalized, dims except channel)
dX     = dOut * scale / sqrt(var + eps)

inputs must contain exactly one tensor (X, matching the Forward call). Forward must have been called first to populate cached intermediates.

func (*BatchNormalization[T]) Forward added in v0.2.1 

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

Forward computes batch normalization in inference mode. inputs: [X, scale, B, mean, var].

func (*BatchNormalization[T]) OpType added in v0.2.1 

func (b *BatchNormalization[T]) OpType() string

OpType returns "BatchNormalization".

func (*BatchNormalization[T]) OutputShape added in v0.2.1 

func (b *BatchNormalization[T]) OutputShape() []int

OutputShape returns the output shape (populated after Forward).

func (*BatchNormalization[T]) Parameters added in v0.2.1 

func (b *BatchNormalization[T]) Parameters() []*graph.Parameter[T]

Parameters returns the learnable parameters (scale and bias) when set, or nil for inference-only mode.

func (*BatchNormalization[T]) SetSaver added in v1.49.0 

func (b *BatchNormalization[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware.

type GroupNormalization added in v1.54.0 

type GroupNormalization[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

GroupNormalization implements Group Normalization (Wu & He, 2018) -- the canonical normalization of convolutional VAE/UNet/diffusion decoders (E127 video VAE, T127.4.1). For input X [N, C, *spatial] it splits the C channels into `numGroups` groups and normalizes each group's (C/groups channels x all spatial positions) jointly, per sample, then applies a learnable per-channel affine (scale, bias of shape [C]): 

y = scale[c] * (x - mean_g) / sqrt(var_g + epsilon) + bias[c]

where mean_g/var_g are computed over the C/groups channels and every spatial position within group g. All math flows through Engine primitives so it appears in the ExecutionPlan tape and runs on CPU/GPU. groups must divide C.

When scale/bias are graph.Parameters (NewGroupNormalizationWithParams), Backward accumulates their gradients and returns dX. The group statistics fall out of the same reduce/elementwise ops as LayerNorm, applied to the grouped [N*groups, (C/groups)*S] view.

func NewGroupNormalization added in v1.54.0 

func NewGroupNormalization[T tensor.Numeric](engine compute.Engine[T], ops numeric.Arithmetic[T], numGroups int, epsilon T) *GroupNormalization[T]

NewGroupNormalization creates an inference layer with no affine (scale/bias nil).

func NewGroupNormalizationWithParams added in v1.54.0 

func NewGroupNormalizationWithParams[T tensor.Numeric](engine compute.Engine[T], ops numeric.Arithmetic[T], numGroups int, epsilon T, scale, bias *graph.Parameter[T]) *GroupNormalization[T]

NewGroupNormalizationWithParams creates a layer with learnable per-channel scale and bias (each shape [C]).

func (*GroupNormalization[T]) Attributes added in v1.54.0 

func (g *GroupNormalization[T]) Attributes() map[string]interface{}

Attributes returns the layer configuration.

func (*GroupNormalization[T]) Backward added in v1.54.0 

func (g *GroupNormalization[T]) Backward(ctx context.Context, _ types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward returns dX and accumulates scale/bias gradients. inputs: [X].

func (*GroupNormalization[T]) Forward added in v1.54.0 

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

Forward computes group normalization. inputs: [X].

func (*GroupNormalization[T]) OpType added in v1.54.0 

func (g *GroupNormalization[T]) OpType() string

OpType returns "GroupNormalization".

func (*GroupNormalization[T]) OutputShape added in v1.54.0 

func (g *GroupNormalization[T]) OutputShape() []int

OutputShape returns the output shape (same as input; populated after Forward).

func (*GroupNormalization[T]) Parameters added in v1.54.0 

func (g *GroupNormalization[T]) Parameters() []*graph.Parameter[T]

Parameters returns the learnable affine parameters when set.

func (*GroupNormalization[T]) SetSaver added in v1.54.0 

func (g *GroupNormalization[T]) SetSaver(sv graph.Saver[T])

SetSaver implements graph.SaverAware.

type LayerNormalization 

type LayerNormalization[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

LayerNormalization implements the Layer Normalization operation.

When T is float32 or sub-float32 (float16, float8) AND the engine is a CPU engine, Backward runs all per-element arithmetic in float64 and casts the result back to T. Parameter storage and gradient outputs stay at T. This removes the catastrophic-cancellation risk in (input - mean) when activations drift far from zero, and protects the stdDev^3 division chain from float32 rounding noise. The GPU engine path preserves the original T-only computation because GPU kernels batch these ops and maintain intermediate precision better than naive per-element float32 math; upgrading GPU kernels to mixed precision is a follow-up once native kernels exist.

func NewLayerNormalization 

func NewLayerNormalization[T tensor.Numeric](engine compute.Engine[T], featureDim int, options ...LayerNormalizationOption[T]) (*LayerNormalization[T], error)

NewLayerNormalization creates a new LayerNormalization layer. featureDim: The dimension over which to normalize (typically the last dimension).

func NewLayerNormalizationFromParams added in v1.9.0 

func NewLayerNormalizationFromParams[T tensor.Numeric](
	engine compute.Engine[T],
	epsilon T,
	gamma *graph.Parameter[T],
	beta *graph.Parameter[T],
) *LayerNormalization[T]

NewLayerNormalizationFromParams creates a LayerNormalization layer from existing gamma (weight) and beta (bias) parameters. This is used for constructing layers from pre-loaded GGUF tensors during model loading.

func (*LayerNormalization[T]) Attributes added in v0.2.1 

func (ln *LayerNormalization[T]) Attributes() map[string]interface{}

Attributes returns the attributes of the LayerNormalization layer.

func (*LayerNormalization[T]) Backward 

func (ln *LayerNormalization[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes the gradients for LayerNormalization.

func (*LayerNormalization[T]) Forward 

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

Forward computes the Layer Normalization.

func (*LayerNormalization[T]) OpType added in v0.2.1 

func (ln *LayerNormalization[T]) OpType() string

OpType returns the operation type of the LayerNormalization layer.

func (*LayerNormalization[T]) OutputShape 

func (ln *LayerNormalization[T]) OutputShape() []int

OutputShape returns the output shape, which is the same as the input shape.

func (*LayerNormalization[T]) Parameters 

func (ln *LayerNormalization[T]) Parameters() []*graph.Parameter[T]

Parameters returns the trainable gamma and beta parameters.

type LayerNormalizationOption added in v0.2.0 

type LayerNormalizationOption[T tensor.Numeric] func(*LayerNormalizationOptions[T])

LayerNormalizationOption is a functional option for configuring LayerNormalization layers.

func WithLayerNormEngineBackward added in v1.50.0 

func WithLayerNormEngineBackward[T tensor.Numeric]() LayerNormalizationOption[T]

WithLayerNormEngineBackward forces Backward to run as engine ops in the element type T, overriding the host float64 mixed-precision backward that shouldUseMixedPrecisionBackward enables for low-precision types.

The mixed-precision backward reads the input, upstream gradient, and gamma back to the host every step -- D2H copies plus host math that are illegal inside a CUDA-graph capture region (compute.GraphCapturer / training.CaptureReplayRunner). Callers that capture training steps must opt in to the engine path and supply their own overflow guards (gradient clipping, divergence detection): the float64 backward exists because float32 LayerNorm backward can overflow on early-training gradient spikes (the CrossAsset f32 cliff).

Gradient storage semantics also change: gamma/beta gradients are ASSIGNED as fresh engine tensors each Backward (the Linear/Bias pattern) instead of host-accumulated via Parameter.AddGradient. Trainers built on DefaultBackpropStrategy accumulate them across steps via the persistent-gradient hook (issue #850); bespoke trainers that relied on AddGradient's in-place accumulation must not use this option.

func WithLayerNormEpsilon added in v0.2.0 

func WithLayerNormEpsilon[T tensor.Numeric](epsilon T) LayerNormalizationOption[T]

WithLayerNormEpsilon sets the epsilon parameter for LayerNormalization.

type LayerNormalizationOptions added in v0.2.0 

type LayerNormalizationOptions[T tensor.Numeric] struct {
	Epsilon T // Small constant to avoid division by zero
	// EngineBackward forces the engine-ops backward path; see
	// WithLayerNormEngineBackward.
	EngineBackward bool
}

LayerNormalizationOptions holds configuration options for LayerNormalization layers.

type RMSNorm added in v0.2.0 

type RMSNorm[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

RMSNorm is a struct that implements the graph.Node interface for RMSNorm.

func NewRMSNorm added in v0.2.0 

func NewRMSNorm[T tensor.Numeric](name string, engine compute.Engine[T], ops numeric.Arithmetic[T], modelDim int, options ...RMSNormOption[T]) (*RMSNorm[T], error)

NewRMSNorm creates a new RMSNorm layer.

func NewRMSNormFromParam added in v0.2.0 

func NewRMSNormFromParam[T tensor.Numeric](engine compute.Engine[T], ops numeric.Arithmetic[T], epsilon T, gain *graph.Parameter[T]) (*RMSNorm[T], error)

NewRMSNormFromParam creates a new RMSNorm layer from an existing gain parameter.

func (*RMSNorm[T]) Attributes added in v0.2.0 

func (r *RMSNorm[T]) Attributes() map[string]interface{}

Attributes returns the attributes.

func (*RMSNorm[T]) Backward added in v0.2.0 

func (r *RMSNorm[T]) Backward(ctx context.Context, mode types.BackwardMode, dOut *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward computes the backward pass of the RMSNorm layer.

func (*RMSNorm[T]) Forward added in v0.2.0 

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

Forward computes the forward pass of the RMSNorm layer.

func (*RMSNorm[T]) OpType added in v0.2.0 

func (r *RMSNorm[T]) OpType() string

OpType returns the operation type.

func (*RMSNorm[T]) OutputShape added in v0.2.0 

func (r *RMSNorm[T]) OutputShape() []int

OutputShape returns the output shape of the RMSNorm layer.

func (*RMSNorm[T]) Parameters added in v0.2.0 

func (r *RMSNorm[T]) Parameters() []*graph.Parameter[T]

Parameters returns the learnable parameters of the RMSNorm layer.

func (*RMSNorm[T]) SetName added in v0.2.0 

func (r *RMSNorm[T]) SetName(name string)

SetName sets the name of the RMSNorm layer.

type RMSNormOption added in v0.2.0 

type RMSNormOption[T tensor.Numeric] func(*RMSNormOptions[T])

RMSNormOption is a functional option for configuring RMSNorm layers.

func WithRMSNormEpsilon added in v0.2.0 

func WithRMSNormEpsilon[T tensor.Numeric](epsilon T) RMSNormOption[T]

WithRMSNormEpsilon sets the epsilon parameter for RMSNorm.

type RMSNormOptions added in v0.2.0 

type RMSNormOptions[T tensor.Numeric] struct {
	Epsilon T // Small constant to avoid division by zero
}

RMSNormOptions holds configuration options for RMSNorm layers.

type SimplifiedLayerNormalization added in v0.2.0 

type SimplifiedLayerNormalization[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

SimplifiedLayerNormalization implements a simplified version of layer normalization.

func NewSimplifiedLayerNormalization added in v0.2.0 

func NewSimplifiedLayerNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	gain *tensor.TensorNumeric[T],
	epsilon T,
) (*SimplifiedLayerNormalization[T], error)

NewSimplifiedLayerNormalization creates a new SimplifiedLayerNormalization layer.

func (*SimplifiedLayerNormalization[T]) Attributes added in v0.2.0 

func (sln *SimplifiedLayerNormalization[T]) Attributes() map[string]interface{}

Attributes returns the attributes of the SimplifiedLayerNormalization layer.

func (*SimplifiedLayerNormalization[T]) Backward added in v0.2.0 

func (sln *SimplifiedLayerNormalization[T]) Backward(ctx context.Context, mode types.BackwardMode, outputGradient *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward applies the backward pass of the SimplifiedLayerNormalization layer.

func (*SimplifiedLayerNormalization[T]) Forward added in v0.2.0 

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

Forward applies the forward pass of the SimplifiedLayerNormalization layer.

func (*SimplifiedLayerNormalization[T]) OpType added in v0.2.0 

func (sln *SimplifiedLayerNormalization[T]) OpType() string

OpType returns the operation type of the SimplifiedLayerNormalization layer.

func (*SimplifiedLayerNormalization[T]) OutputShape added in v0.2.0 

func (sln *SimplifiedLayerNormalization[T]) OutputShape() []int

OutputShape returns the output shape of the layer.

func (*SimplifiedLayerNormalization[T]) Parameters added in v0.2.0 

func (sln *SimplifiedLayerNormalization[T]) Parameters() []*graph.Parameter[T]

Parameters returns the learnable parameters of the layer.

type SkipSimplifiedLayerNormalization added in v0.2.0 

type SkipSimplifiedLayerNormalization[T tensor.Numeric] struct {
	// contains filtered or unexported fields
}

SkipSimplifiedLayerNormalization applies SimplifiedLayerNormalization and adds a residual connection.

func NewSkipSimplifiedLayerNormalization added in v0.2.0 

func NewSkipSimplifiedLayerNormalization[T tensor.Numeric](
	engine compute.Engine[T],
	ops numeric.Arithmetic[T],
	gain *tensor.TensorNumeric[T],
	epsilon T,
) (*SkipSimplifiedLayerNormalization[T], error)

NewSkipSimplifiedLayerNormalization creates a new SkipSimplifiedLayerNormalization layer.

func (*SkipSimplifiedLayerNormalization[T]) Attributes added in v0.2.0 

func (sln *SkipSimplifiedLayerNormalization[T]) Attributes() map[string]interface{}

Attributes returns the attributes of the underlying normalization layer.

func (*SkipSimplifiedLayerNormalization[T]) Backward added in v0.2.0 

func (sln *SkipSimplifiedLayerNormalization[T]) Backward(ctx context.Context, mode types.BackwardMode, outputGrad *tensor.TensorNumeric[T], inputs ...*tensor.TensorNumeric[T]) ([]*tensor.TensorNumeric[T], error)

Backward applies the backward pass of the SkipSimplifiedLayerNormalization layer.

func (*SkipSimplifiedLayerNormalization[T]) Forward added in v0.2.0 

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

Forward applies the forward pass of the SkipSimplifiedLayerNormalization layer.

func (*SkipSimplifiedLayerNormalization[T]) OpType added in v0.2.0 

func (sln *SkipSimplifiedLayerNormalization[T]) OpType() string

OpType returns the operation type of the SkipSimplifiedLayerNormalization layer.

func (*SkipSimplifiedLayerNormalization[T]) OutputShape added in v0.2.0 

func (sln *SkipSimplifiedLayerNormalization[T]) OutputShape() []int

OutputShape returns the output shape of the layer.

func (*SkipSimplifiedLayerNormalization[T]) Parameters added in v0.2.0 

func (sln *SkipSimplifiedLayerNormalization[T]) Parameters() []*graph.Parameter[T]

Parameters returns the learnable parameters of the layer.

Source Files

View all Source files

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.