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 trait

ReduceOp

Trait for an associative reduction monoid (identity + associative join).

Required: __init__ (identity), accumulate[w] (SIMD-tile fold), join (sequential combine). Optional: join_parallel[R: Reducer] (parallel combine across participants). The trait's default delegates to reducer.generic (only needs join); override to use a Reducer's hardware-fast scalar primitives, or a field-wise algebraic decomposition that collapses many joins into a few scalar reductions.

TrivialRegisterPassable because accumulators pass by value through the parallel reducer.

accumulate carries a free val_dtype parameter because Mojo trait conformance treats a struct-param dtype and a trait-declared comptime dtype as distinct types; impls cast val to the accumulator dtype inline.

Constraints: join must be associative. The default join_parallel also assumes commutativity (participants reduce in unspecified order); non-commutative monoids must override it.

Implemented traits​

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

comptime members​

Single​

comptime Single

The width-1 form of this monoid (Self at W == 1), produced by reduce and consumed by the scalar broadcast constructor.

width​

comptime width

The SIMD width W of the state, exposed so the trait's default reduce / broadcast iterate the lanes generically.

Required methods​

__init__​

def __init__() -> Self

Returns the identity element.

Returns:

_Self

def __init__(out self, *, copy: Self)

Create a new instance of the value by copying an existing one.

Args:

  • ​copy (_Self): The value to copy.

Returns:

_Self

def __init__(out self, *, deinit move: Self)

Create a new instance of the value by moving the value of another.

Args:

  • ​move (_Self): The value to move.

Returns:

_Self

__getitem__​

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

Returns lane j of the state as a width-1 monoid.

Returns:

_Self.Single

__setitem__​

def __setitem__(mut self, j: Int, s: _Self.Single)

Writes the width-1 monoid s into lane j of the state.

accumulate​

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

Folds a SIMD tile into self elementwise — lane j folds (val[j], idx[j]) — casting to the accumulator dtype internally.

idx[j] is where val[j] sits along the reduce axis. The scaffolder builds it per tier — consecutive positions (base + iota) when the lanes are partials of one output, or a constant (base) when the lanes are independent output columns — so the monoid never learns its tier. Value-only monoids (ReduceSum, ReduceMax, ReduceMin, ReduceProduct, OnlineLogSumExp) ignore idx; index-tracking monoids (ArgMax, ArgMin) fold it directly. Defaults to a zero vector for callers that don't track positions.

Parameters:

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

Args:

join​

def join(mut self, other: Self)

Combines self with other associatively (sequential).

Args:

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

Provided methods​

__init__​

def __init__(s: _Self.Single) -> Self

Broadcasts a width-1 result across all W lanes — the inverse of reduce, so the body can read result.slice[w] uniformly. Default: write every lane via __setitem__.

Returns:

_Self

reduce​

def reduce(self) -> _Self.Single

Collapses the W lane-wise partials of one output into a single width-1 result.

The scaffolder calls it once, immediately before join_parallel, and only on cooperative tiers (W lanes are partials of one output). On the tiled tier the W lanes are independent outputs, so both reduce and join_parallel are skipped and the per-lane results stand.

Default: fold the lanes through join (the loop is empty at W == 1, so this is a branchless identity there). Override for a faster horizontal reduce (e.g. acc.reduce_add() for ReduceSum).

Returns:

_Self.Single

join_parallel​

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

Combines self across all participants via reducer, leaving the combined value on every participant (and, for SIMD states, every lane).

Runs after reduce, so self's lanes already hold a single within-participant result; this only does the cross-participant step. Default: reducer.generic(self) — correct for any commutative + associative monoid, and since reduce left all lanes equal, the generic field-wise combine yields the result on every lane (no separate broadcast). Override to dispatch to reducer.sum/max/min for a faster per-dtype path, splatting the scalar result back across lanes.

Parameters:

  • ​R (Reducer): The parallel scalar reducer.

Args:

  • ​reducer (R): The reducer instance.

pad​

static def pad[Wout: Int, Tout: DType, identity: Scalar[Tout], val_dtype: DType, w_in: Int](val: SIMD[val_dtype, w_in]) -> SIMD[Tout, Wout]

Pads val to width Wout, casting val_dtype -> Tout and filling lanes [w_in, Wout) with identity.

Lets a monoid's accumulate[w] body stay a single expression — no comptime if w == Self.W branching. For w_in == Wout returns the cast directly; for w_in < Wout builds an identity-filled vector and overwrites lanes [0, w_in) with the cast val.

Precondition: w_in <= Wout. Bodies arrange this by computing W = rowwise.pick_simd_width[...] once across all monoids and dtypes involved, so the scaffolder's tile width never exceeds any monoid's Self.W.

The default suffices for single-field SIMD monoids. Multi-field monoids (ArgMax, ArgMin) call pad twice — once per field — with each field's own identity.

Parameters:

  • ​Wout (Int): Output SIMD width.
  • ​Tout (DType): Output element dtype.
  • ​identity (Scalar[Tout]): Padding value for tail lanes.
  • ​val_dtype (DType): Input element dtype.
  • ​w_in (Int): Input SIMD width.

Args:

Returns:

SIMD[Tout, Wout]: A SIMD[Tout, Wout] with val in lanes [0, w_in) and identity in lanes [w_in, Wout).

copy​

def copy(self) -> Self

Explicitly construct a copy of self, a convenience method for Self(copy=self) when the type is inconvenient to write out.

Overriding this method is not allowed.

Returns:

_Self: A copy of this value.

Models

Cookbook

Puzzles

Blog

Community

Contact