IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

Welford

struct Welford[dtype: DType, W: Int = Int(1)]

Welford's online mean/variance monoid (Welford 1962, Chan et al. 1979 for the combine) β€” the reduction half of layer_norm / group_norm.

Holds {count, mean, M2} β€” running sample count, mean, and sum of squared deviations from the mean β€” lane-wise across a width-W SIMD field. Variance is M2 / count. Combining two sub-states uses Chan's formula:

n   = na + nb
Ξ΄   = ΞΌb - ΞΌa
ΞΌ   = ΞΌa + Ξ΄ * nb / n
M2  = M2_a + M2_b + δ² * na * nb / n

accumulate[W] is elementwise: each call folds one new sample into every lane's Welford (Chan with nb = 1). Because every lane advances its count in lock-step, 1 / count is a single scalar reciprocal β€” sidestepping the per-lane SIMD fdiv that serializes through the GPU divider β€” with no per-tile horizontal reduce. reduce Chan-combines the W lanes once at the end. On the tiled tier the lanes are independent output columns, so reduce/join_parallel are skipped and each lane's Welford stands as that column's statistics.

Triple-field state: count, mean, M2, each a SIMD[dtype, W]. join_parallel cross-thread-combines through reducer.generic β€” Welford's combine isn't a hardware primitive.

Parameters​

  • ​dtype (DType): The accumulator dtype (must be floating-point).
  • ​W (Int): SIMD width of the per-field storage. Defaults to 1 (scalar state). Cooperative tiers pass the body's SIMD width so each lane accumulates a strided slice of the row; the tiled tier passes tile_width so each lane is an output column.

Fields​

  • ​count (SIMD[dtype, W]): Per-lane running sample count (uniform across lanes during accumulation). After reduce, all lanes hold the combined count.
  • ​mean (SIMD[dtype, W]): Per-lane running mean. After reduce, all lanes hold the combined mean (cooperative); on the tiled tier each lane is its column's mean.
  • ​M2 (SIMD[dtype, W]): Per-lane running sum-of-squared-deviations-from-the-mean. Variance is M2 / count. After reduce, all lanes hold the combined M2 (cooperative); on the tiled tier each lane is its column's M2.

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, ReduceOp, RegisterPassable, TrivialRegisterPassable

comptime members​

Single​

comptime Single = Welford[dtype]

width​

comptime width = W

Methods​

__init__​

def __init__() -> Self

Identity: every lane at (0, 0, 0).

__getitem__​

def __getitem__(self, j: Int) -> Self.Single

Returns lane j as a width-1 monoid.

Returns:

Self.Single

__setitem__​

def __setitem__(mut self, j: Int, s: Welford[dtype])

Writes width-1 monoid s into lane j.

accumulate​

def accumulate[val_dtype: DType, w: Int](mut self, val: SIMD[val_dtype, w], idx: SIMD[DType.int64, w] = 0)

Elementwise Welford fold: each lane folds one new sample into its running state (Chan with nb = 1).

For a full tile (w == W) every lane advances its count in lock-step, so 1 / new_count is a single scalar reciprocal and the SIMD work is one sub + two fmas β€” no horizontal reduce, no per-lane divide. A partial tile (w < W, the scalar tail of an odd row) masks the count increment to its first w lanes, which costs a per-lane reciprocal; tiers at W > 1 only hit full tiles (the row divides the SIMD width), so this path is compiled but never executed there.

Parameters:

  • ​val_dtype (DType): The tile's source dtype.
  • ​w (Int): SIMD width of the tile.

Args:

join​

def join(mut self, other: Self)

Sequential combine of two Welford states via Chan's formula, lane-wise.

Args:

  • ​other (Self): The state to combine into self.

join_parallel​

def join_parallel[R: Reducer](mut self, reducer: R)

Cross-thread combine via reducer.generic (Welford's combine isn't a hardware primitive), then broadcast the result from lane 0 across all lanes so a body reads (mean, M2, count).slice[w] uniformly. Runs after reduce.

Parameters:

  • ​R (Reducer): The parallel reducer.

Args:

  • ​reducer (R): The reducer instance.

Models

Cookbook

Puzzles

Blog

Community

Contact