Struct SegmentMeanPlan 

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

segment_mean plan (sorted).

out[s, d] = mean_{n : segment_ids[n] == s} input[n, d] (TF / JAX segment_mean). Requires segment_ids monotonically non-decreasing.

When to use: forward sorted segment-mean. Pair with SegmentMeanBackwardPlan.

Dtypes: {f32, f64}.

Shape limits: input [N, D]; segment_ids [N] with values in [0, num_segments); output [num_segments, D].

Workspace: none — segment counts derived inline via binary search.

Precision guarantee: deterministic, bit-stable.

Index policy: out-of-range IDs dropped. Empty segments emit zero (division is guarded; no NaN).

Implementations

impl<T: Element> SegmentMeanPlan<T>

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

Pick a kernel.

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

Validate args.

pub fn workspace_size(&self) -> usize

Workspace size — zero (count computed inline via binary search).

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: SegmentMeanArgs<'_, T>, ) -> Result<()>

Launch.