Struct ArgmaxTracker 

pub struct ArgmaxTracker { /* private fields */ }

Tracks winner indices from tropical forward-pass operations.

During a tropical contraction C[i,j] = max_k (A[i,k] + B[k,j]), the tracker records which k achieved the maximum for each (i,j). The backward pass uses these indices to route gradients.

Examples

use tenferro_ext_tropical::ArgmaxTracker;

// Create a tracker for a 3x5 output
let tracker = ArgmaxTracker::new(&[3, 5]);

// After forward pass, query the winner index for output element (1, 2)
let k_winner = tracker.winner_index(&[1, 2]);
assert_eq!(k_winner, 0); // initialized to 0

Implementations

impl ArgmaxTracker

pub fn new(output_shape: &[usize]) -> Self

Create a new tracker for an output of the given shape.

All winner indices are initialized to 0.

Examples

use tenferro_ext_tropical::ArgmaxTracker;

let tracker = ArgmaxTracker::new(&[3, 5]);
assert_eq!(tracker.output_shape(), &[3, 5]);

pub fn output_shape(&self) -> &[usize]

Return the output shape.

pub fn indices(&self) -> &[usize]

Return the winner indices as a flat slice.

pub fn indices_mut(&mut self) -> &mut [usize]

Return a mutable reference to the winner indices.

pub fn winner_index(&self, position: &[usize]) -> usize

Look up the winner index for a given multi-dimensional output position.

Examples

use tenferro_ext_tropical::ArgmaxTracker;

let tracker = ArgmaxTracker::new(&[3, 5]);
let k = tracker.winner_index(&[1, 2]);
assert_eq!(k, 0); // initialized to 0

Panics

Panics if position has the wrong rank or any index is out of bounds for the tracked output shape.