Struct SegmentSumArgs 

pub struct SegmentSumArgs<'a, T: Element> {
    pub input: TensorRef<'a, T, 2>,
    pub segment_ids: TensorRef<'a, i32, 1>,
    pub output: TensorMut<'a, T, 2>,
}

Args bundle for a segment_sum launch.

Fields

input: TensorRef<'a, T, 2>

Input [N, D].

segment_ids: TensorRef<'a, i32, 1>

Segment ids [N], i32, sorted non-decreasing, values in [0, num_segments).

output: TensorMut<'a, T, 2>

Output [num_segments, D]. Overwritten by the launch — no accumulation into pre-existing state.