Struct SegmentMaxArgs 

pub struct SegmentMaxArgs<'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_max launch.

Fields

input: TensorRef<'a, T, 2>

Input [N, D].

segment_ids: TensorRef<'a, i32, 1>

Segment ids [N] — sorted non-decreasing.

output: TensorMut<'a, T, 2>

Output [num_segments, D].