Struct SegmentMaxPlan 

pub struct SegmentMaxPlan<T: Element> { /* private fields */ }

segment_max plan (sorted, FW only).

out[s, d] = max_{n : segment_ids[n] == s} input[n, d] (TF / JAX segment_max). Requires sorted segment_ids.

When to use: forward sorted segment-max. No BW plan — argmax tracking from FW is needed (deferred).

Dtypes: {f32, f64}.

Shape limits: input [N, D]; segment_ids [N]; output [num_segments, D].

Workspace: none.

Precision guarantee: deterministic, bit-stable.

Index policy: out-of-range IDs dropped. Empty segments emit zero (per-op identity sentinel from segment_sorted_kernel).

Implementations

impl<T: Element> SegmentMaxPlan<T>

pub fn select( _stream: &Stream, desc: &SegmentMaxDescriptor, _pref: PlanPreference, ) -> Result<Self>

Pick a kernel.

pub fn can_implement(&self, args: &SegmentMaxArgs<'_, T>) -> Result<()>

Validate args.

pub fn workspace_size(&self) -> usize

Workspace size — zero.

pub fn sku(&self) -> KernelSku

Identity of the kernel this plan picked.

pub fn precision_guarantee(&self) -> PrecisionGuarantee

Numerical guarantees.

pub fn run( &self, stream: &Stream, _workspace: Workspace<'_>, args: SegmentMaxArgs<'_, T>, ) -> Result<()>

Launch.