rlx_ir/ops/
reduction.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Reduction builders: reduce, softmax, cumsum, sample
6//! (plan #53).
7
8use crate::op::ReduceOp;
9use crate::{Graph, NodeId, Op, Shape};
10
11impl Graph {
12    /// Reduce.
13    pub fn reduce(
14        &mut self,
15        input: NodeId,
16        op: ReduceOp,
17        axes: Vec<usize>,
18        keep_dim: bool,
19        shape: Shape,
20    ) -> NodeId {
21        self.push(Op::Reduce { op, axes, keep_dim }, vec![input], shape, None)
22    }
23
24    /// Softmax.
25    pub fn softmax(&mut self, input: NodeId, axis: i32, shape: Shape) -> NodeId {
26        self.push(Op::Softmax { axis }, vec![input], shape, None)
27    }
28
29    /// Histogram of `input` into `bins` equal-width buckets over `[min, max]`.
30    /// Output is a 1-D f32 tensor `[bins]` of counts. Elements outside the
31    /// range are dropped and `x == max` lands in the last bin (matches
32    /// `numpy.histogram`). Non-differentiable.
33    pub fn histogram(&mut self, input: NodeId, bins: usize, min: f32, max: f32) -> NodeId {
34        let out = Shape::new(&[bins], crate::DType::F32);
35        self.push(Op::Histogram { bins, min, max }, vec![input], out, None)
36    }
37
38    /// Cumulative sum along an axis (output shape == input shape).
39    pub fn cumsum(&mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape) -> NodeId {
40        self.push(Op::Cumsum { axis, exclusive }, vec![input], shape, None)
41    }
42
43    /// Cumulative product along an axis (output shape == input shape).
44    pub fn cumprod(&mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape) -> NodeId {
45        self.push(Op::CumProd { axis, exclusive }, vec![input], shape, None)
46    }
47
48    /// Cumulative maximum along an axis (output shape == input shape).
49    pub fn cummax(&mut self, input: NodeId, axis: i32, exclusive: bool, shape: Shape) -> NodeId {
50        self.push(Op::CumMax { axis, exclusive }, vec![input], shape, None)
51    }
52
53    /// Index of the max along `axis` (f32-encoded indices).
54    pub fn argmax(&mut self, input: NodeId, axis: usize, keep_dim: bool, shape: Shape) -> NodeId {
55        self.push(Op::ArgMax { axis, keep_dim }, vec![input], shape, None)
56    }
57
58    /// Index of the min along `axis` (f32-encoded indices).
59    pub fn argmin(&mut self, input: NodeId, axis: usize, keep_dim: bool, shape: Shape) -> NodeId {
60        self.push(Op::ArgMin { axis, keep_dim }, vec![input], shape, None)
61    }
62
63    /// Fused sample: logits → token id (one f32-encoded id per row).
64    /// `output_shape` should be `[batch]` (one id per logit row).
65    pub fn sample(
66        &mut self,
67        logits: NodeId,
68        top_k: usize,
69        top_p: f32,
70        temperature: f32,
71        seed: u64,
72        output_shape: Shape,
73    ) -> NodeId {
74        self.push(
75            Op::Sample {
76                top_k,
77                top_p,
78                temperature,
79                seed,
80            },
81            vec![logits],
82            output_shape,
83            None,
84        )
85    }
86}
rlx_ir/ops/reduction.rs

rlx_ir/ops/
reduction.rs
