Struct UnsortedSegmentProdArgs 

pub struct UnsortedSegmentProdArgs<'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 an unsorted_segment_prod launch.

Fields

input: TensorRef<'a, T, 2>

Input [N, D].

segment_ids: TensorRef<'a, i32, 1>

Segment ids [N], i32, in any order.

output: TensorMut<'a, T, 2>

Output [num_segments, D]. Overwritten by the launch — kernel fills 1.0 before the atomic accumulation phase.