Struct OnlineSoftmax 

pub struct OnlineSoftmax<E: Float> { /* private fields */ }

The running max m and sum l of every slice along axis of a score, in the registers of its holder.

The score is cut where its holder keeps it (AxisSlices): a unit’s register block holds whole slices of its own, and a plane’s window of shared memory is the plane’s, every unit of it keeping the state of every slice.

Implementations

impl<E: Float> OnlineSoftmax<E>

pub fn along(score: &Tile, axis: Axis) -> OnlineSoftmax<E>

The state of every slice along axis of a score held as score is, before any block is taken in: each unit keeps the slices it holds cells of (Tile::slices_per_unit).

pub fn new(slices: usize, axis: Axis) -> OnlineSoftmax<E>

The state of slices slices along axis before any block is taken in: what a holder opens before the walk that hands it its score, as a cube’s plane does before the cube’s ring.

pub fn step(&mut self, score: &Tile, bias: &Tile, scale: E) -> Array<E>

Take one block of score into the slices. The block becomes score · scale + bias, the slices’ max and sum move on, and score is left holding exp(score − max): the block’s probabilities, unnormalized. Returns each slice’s exp(max before − max after), what anything summed against the earlier blocks is multiplied by (AxisSlices::mul).

bias is a procedural tile windowed to the same region as score. A masked cell’s bias is E::min_value(); a slice whose every cell so far is masked holds zeros and a zero sum. On a plane’s window every unit of the plane calls it, and the caller meets the plane before it and before the probabilities are read.

pub fn normalizer(&mut self, plane: &Region) -> Array<E>

What this plane’s sum against the values is multiplied by once its last block is taken in: each slice’s 1 / l, or, where the cube’s planes split this axis, each holding part of the same slices, the correction to the merged max times 1 / L of the merged sum, so the planes’ products only remain to be added (PlanesOutput). plane is this plane’s region of the cube; the partitioning it was cut by says how many planes share its slices.

Where planes split the axis it is a meeting of the cube: every unit of every plane calls it, each holding its plane’s state as a plane’s window or fragments keep it, and it meets the cube three times through shared memory. The planes’ values are read in plane order, so the merged sum is the same bits from run to run. A slice every plane masked stays at zero.

pub fn relayed( &mut self, relay: &Relay<'_>, carried_max: &mut Tile, carried_sum: &mut Tile, ) -> RelayedFactors<E>

What this plane’s sum and the carry’s rows are multiplied by in its cube’s turn at relay, where the cubes of the relay split this axis, each holding part of the same slices: the cross-cube normalizer, the carry taking the turns’ partials one at a time.

carried_max and carried_sum are this plane’s window of the slices’ state the relay holds beside its carry: the max and sum every turn before this one merged. This plane merges its own into them and leaves the result there for the next turn. Each factor is a slice’s exp(max − merged max), of this plane’s state for its sum and of the carried one for the carry, both times 1 / L of the merged sum on the last turn, so the carry that turn hands on is the answer. The first turn reads nothing, its drain replacing the carry. A slice every turn masked stays at zero.

Every unit of every plane calls it once its cube has taken the turn; the first unit of each plane writes the state, which the relay publishes with the carry. The caller rescales the carry (Relay::carried) and its sum by the factors, then drains.

pub fn state(&self) -> OnlineSoftmaxState<E>

Every slice’s running max and sum as they stand, copied out: nothing is divided and the state walks on unchanged.

pub fn store_rows(&self, score: &Tile, max: &mut Tile, sum: &mut Tile)

Every slice’s running max and sum as they stand, written into max and sum, a plane’s windows of one cell a slice in slice order: what a holder leaves beside its undivided sum where its slices merge with other holders’ parts of them, and the one way to read out a state taken in from a manual-mma grid, whose units each keep the rows they hold cells of.

score is the tile the state was taken in from (along), which says which unit holds which row: each writes its own, one cell a row. A state every unit holds whole is written by the plane’s first unit. Every unit of the plane calls it, and the caller meets the plane before the windows are read.

pub fn store_logsumexp(&self, score: &Tile, logsumexp: &mut Tile)

Every slice’s logsumexp, m + ln(l), written into logsumexp, a window of one cell a slice in slice order that a unit or a plane holds: what a pass that walks the slices again reads its probabilities back from (AxisSlices::probabilities). A slice whose every cell was masked keeps the minimum value, which reads back as probabilities of zero.

score is the tile the state was taken in from, which says which unit holds which slice, as store_rows reads it. Every unit of the holder calls it.

pub fn recip_l(&self) -> Array<E>

Each slice’s 1 / l, what its sum of probabilities times values is normalized by; zero for a slice whose every cell was masked.

pub fn __expand_along( scope: &Scope, score: &<Tile as CubeType>::ExpandType, axis: Axis, ) -> <OnlineSoftmax<E> as CubeType>::ExpandType

The state of every slice along axis of a score held as score is, before any block is taken in: each unit keeps the slices it holds cells of (Tile::slices_per_unit).

pub fn __expand_new( scope: &Scope, slices: usize, axis: Axis, ) -> <OnlineSoftmax<E> as CubeType>::ExpandType

The state of slices slices along axis before any block is taken in: what a holder opens before the walk that hands it its score, as a cube’s plane does before the cube’s ring.

pub fn __expand_step( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, score: &<Tile as CubeType>::ExpandType, bias: &<Tile as CubeType>::ExpandType, scale: <E as CubeType>::ExpandType, ) -> <Array<E> as CubeType>::ExpandType

Take one block of score into the slices. The block becomes score · scale + bias, the slices’ max and sum move on, and score is left holding exp(score − max): the block’s probabilities, unnormalized. Returns each slice’s exp(max before − max after), what anything summed against the earlier blocks is multiplied by (AxisSlices::mul).

bias is a procedural tile windowed to the same region as score. A masked cell’s bias is E::min_value(); a slice whose every cell so far is masked holds zeros and a zero sum. On a plane’s window every unit of the plane calls it, and the caller meets the plane before it and before the probabilities are read.

pub fn __expand_normalizer( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, plane: &<Region as CubeType>::ExpandType, ) -> <Array<E> as CubeType>::ExpandType

What this plane’s sum against the values is multiplied by once its last block is taken in: each slice’s 1 / l, or, where the cube’s planes split this axis, each holding part of the same slices, the correction to the merged max times 1 / L of the merged sum, so the planes’ products only remain to be added (PlanesOutput). plane is this plane’s region of the cube; the partitioning it was cut by says how many planes share its slices.

Where planes split the axis it is a meeting of the cube: every unit of every plane calls it, each holding its plane’s state as a plane’s window or fragments keep it, and it meets the cube three times through shared memory. The planes’ values are read in plane order, so the merged sum is the same bits from run to run. A slice every plane masked stays at zero.

pub fn __expand_relayed( scope: &Scope, this: &mut <Self as CubeType>::ExpandType, relay: &<Relay<'_> as CubeType>::ExpandType, carried_max: &mut <Tile as CubeType>::ExpandType, carried_sum: &mut <Tile as CubeType>::ExpandType, ) -> <RelayedFactors<E> as CubeType>::ExpandType

What this plane’s sum and the carry’s rows are multiplied by in its cube’s turn at relay, where the cubes of the relay split this axis, each holding part of the same slices: the cross-cube normalizer, the carry taking the turns’ partials one at a time.

carried_max and carried_sum are this plane’s window of the slices’ state the relay holds beside its carry: the max and sum every turn before this one merged. This plane merges its own into them and leaves the result there for the next turn. Each factor is a slice’s exp(max − merged max), of this plane’s state for its sum and of the carried one for the carry, both times 1 / L of the merged sum on the last turn, so the carry that turn hands on is the answer. The first turn reads nothing, its drain replacing the carry. A slice every turn masked stays at zero.

Every unit of every plane calls it once its cube has taken the turn; the first unit of each plane writes the state, which the relay publishes with the carry. The caller rescales the carry (Relay::carried) and its sum by the factors, then drains.

pub fn __expand_state( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> <OnlineSoftmaxState<E> as CubeType>::ExpandType

Every slice’s running max and sum as they stand, copied out: nothing is divided and the state walks on unchanged.

pub fn __expand_store_rows( scope: &Scope, this: &<Self as CubeType>::ExpandType, score: &<Tile as CubeType>::ExpandType, max: &mut <Tile as CubeType>::ExpandType, sum: &mut <Tile as CubeType>::ExpandType, )

Every slice’s running max and sum as they stand, written into max and sum, a plane’s windows of one cell a slice in slice order: what a holder leaves beside its undivided sum where its slices merge with other holders’ parts of them, and the one way to read out a state taken in from a manual-mma grid, whose units each keep the rows they hold cells of.

score is the tile the state was taken in from (along), which says which unit holds which row: each writes its own, one cell a row. A state every unit holds whole is written by the plane’s first unit. Every unit of the plane calls it, and the caller meets the plane before the windows are read.

pub fn __expand_store_logsumexp( scope: &Scope, this: &<Self as CubeType>::ExpandType, score: &<Tile as CubeType>::ExpandType, logsumexp: &mut <Tile as CubeType>::ExpandType, )

Every slice’s logsumexp, m + ln(l), written into logsumexp, a window of one cell a slice in slice order that a unit or a plane holds: what a pass that walks the slices again reads its probabilities back from (AxisSlices::probabilities). A slice whose every cell was masked keeps the minimum value, which reads back as probabilities of zero.

score is the tile the state was taken in from, which says which unit holds which slice, as store_rows reads it. Every unit of the holder calls it.

pub fn __expand_recip_l( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> <Array<E> as CubeType>::ExpandType

Each slice’s 1 / l, what its sum of probabilities times values is normalized by; zero for a slice whose every cell was masked.

Trait Implementations

impl<E: Float> CubeType for OnlineSoftmax<E>

type ExpandType = OnlineSoftmaxExpand<E>