rlx_ir/
infer.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Shape-inferred graph builder — ergonomic API that auto-computes output shapes.
17//!
18//! Import [`GraphExt`] and call short-name methods instead of providing explicit shapes:
19//! ```rust
20//! use rlx_ir::*;
21//! use rlx_ir::infer::GraphExt;
22//!
23//! let mut g = Graph::new("example");
24//! let x = g.input("x", Shape::new(&[4, 384], DType::F32));
25//! let w = g.param("w", Shape::new(&[384, 1536], DType::F32));
26//! let b = g.param("b", Shape::new(&[1536], DType::F32));
27//! let mm = g.mm(x, w);
28//! let add = g.add(mm, b);
29//! let out = g.gelu(add);
30//! let two = g.constant(2.0, DType::F32);
31//! let scaled = g.mul(x, two);
32//! let c = g.try_constant(2.0, DType::F32).unwrap(); // fallible variant
33//! g.set_outputs(vec![out, scaled, c]);
34//! ```
35
36use crate::dtype::scalar_constant_bytes;
37use crate::op::*;
38use crate::shape;
39use crate::{DType, Graph, NodeId, Op, Shape};
40
41/// Extension trait for shape-inferred graph building.
42pub trait GraphExt {
43    // ── Linear algebra ──────────────────────────────────────
44    fn mm(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
45
46    // ── Binary ──────────────────────────────────────────────
47    fn add(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
48    fn sub(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
49    fn mul(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
50    fn div(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
51
52    // ── Activation ──────────────────────────────────────────
53    fn gelu(&mut self, x: NodeId) -> NodeId;
54    /// Tanh-approximation GELU (PyTorch's default `gelu` formula,
55    /// also candle's `Tensor::gelu`). Use this when porting models
56    /// whose reference implementations use the tanh form for
57    /// numerical parity (e.g. DINOv2, many ViTs).
58    fn gelu_approx(&mut self, x: NodeId) -> NodeId;
59    fn silu(&mut self, x: NodeId) -> NodeId;
60    fn relu(&mut self, x: NodeId) -> NodeId;
61    fn exp(&mut self, x: NodeId) -> NodeId;
62    fn sqrt(&mut self, x: NodeId) -> NodeId;
63    fn neg(&mut self, x: NodeId) -> NodeId;
64    fn tanh(&mut self, x: NodeId) -> NodeId;
65
66    // ── Normalization ───────────────────────────────────────
67    fn ln(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId;
68    fn layer_norm2d(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId;
69    fn group_norm(
70        &mut self,
71        x: NodeId,
72        gamma: NodeId,
73        beta: NodeId,
74        num_groups: usize,
75        eps: f32,
76    ) -> NodeId;
77    fn rms_norm(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId;
78
79    // ── Convolution (NCHW) ───────────────────────────────────
80    fn conv2d(
81        &mut self,
82        input: NodeId,
83        weight: NodeId,
84        kernel_size: [usize; 2],
85        stride: [usize; 2],
86        padding: [usize; 2],
87        dilation: [usize; 2],
88        groups: usize,
89    ) -> NodeId;
90    fn conv_transpose2d(
91        &mut self,
92        input: NodeId,
93        weight: NodeId,
94        kernel_size: [usize; 2],
95        stride: [usize; 2],
96        padding: [usize; 2],
97        dilation: [usize; 2],
98        output_padding: [usize; 2],
99        groups: usize,
100    ) -> NodeId;
101
102    // ── Reduction ───────────────────────────────────────────
103    fn sum(&mut self, x: NodeId, axes: Vec<usize>, keep_dim: bool) -> NodeId;
104    fn mean(&mut self, x: NodeId, axes: Vec<usize>, keep_dim: bool) -> NodeId;
105    fn sm(&mut self, x: NodeId, axis: i32) -> NodeId;
106
107    // ── Shape manipulation ──────────────────────────────────
108    fn reshape_(&mut self, x: NodeId, new_shape: Vec<i64>) -> NodeId;
109    fn transpose_(&mut self, x: NodeId, perm: Vec<usize>) -> NodeId;
110    fn narrow_(&mut self, x: NodeId, axis: usize, start: usize, len: usize) -> NodeId;
111    fn concat_(&mut self, inputs: Vec<NodeId>, axis: usize) -> NodeId;
112    fn gather_(&mut self, table: NodeId, indices: NodeId, axis: usize) -> NodeId;
113
114    // ── Comparison ──────────────────────────────────────────
115    fn eq(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
116    fn lt(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId;
117
118    // ── Attention ───────────────────────────────────────────
119    fn attention_(
120        &mut self,
121        q: NodeId,
122        k: NodeId,
123        v: NodeId,
124        mask: NodeId,
125        num_heads: usize,
126        head_dim: usize,
127    ) -> NodeId;
128
129    // ── RoPE ────────────────────────────────────────────────
130    fn rope(&mut self, x: NodeId, cos: NodeId, sin: NodeId, head_dim: usize) -> NodeId;
131    /// Partial RoPE: rotate the first `n_rot` dims (NeoX offset `n_rot/2`).
132    fn rope_n(
133        &mut self,
134        x: NodeId,
135        cos: NodeId,
136        sin: NodeId,
137        head_dim: usize,
138        n_rot: usize,
139    ) -> NodeId;
140    /// RoPE with an explicit pairing flavor ([`crate::op::RopeStyle`]).
141    fn rope_styled(
142        &mut self,
143        x: NodeId,
144        cos: NodeId,
145        sin: NodeId,
146        head_dim: usize,
147        style: crate::op::RopeStyle,
148    ) -> NodeId;
149    /// Partial RoPE with an explicit pairing flavor.
150    fn rope_n_styled(
151        &mut self,
152        x: NodeId,
153        cos: NodeId,
154        sin: NodeId,
155        head_dim: usize,
156        n_rot: usize,
157        style: crate::op::RopeStyle,
158    ) -> NodeId;
159
160    // ── Cast ────────────────────────────────────────────────
161    fn cast(&mut self, x: NodeId, to: DType) -> NodeId;
162
163    // ── Literals ────────────────────────────────────────────
164    /// Rank-0 broadcastable scalar (`Op::Constant`). `f16` / `bf16`
165    /// are lowered as `f32` constant + `cast`.
166    fn constant(&mut self, value: f64, dtype: DType) -> NodeId;
167
168    /// Fallible variant of [`GraphExt::constant`]. Returns an error when
169    /// `value` is out of range for `dtype` or when `dtype` cannot be encoded
170    /// directly (callers may lower `f16` / `bf16` via `try_constant` on
171    /// `F32` plus `cast`).
172    fn try_constant(&mut self, value: f64, dtype: DType) -> Result<NodeId, String>;
173
174    // ── Stop gradient ───────────────────────────────────────
175    /// Identity forward, zero backward. The reverse-mode autodiff rule
176    /// for `Op::StopGradient` returns no gradient contribution to the
177    /// input. Equivalent to PyTorch's `tensor.detach()` /
178    /// `jax.lax.stop_gradient` / TF's `tf.stop_gradient`.
179    fn stop_gradient(&mut self, x: NodeId) -> NodeId;
180}
181
182impl GraphExt for Graph {
183    fn mm(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
184        let s =
185            shape::matmul_shape(self.shape(lhs), self.shape(rhs)).expect("matmul shape inference");
186        self.matmul(lhs, rhs, s)
187    }
188
189    fn add(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
190        let s = shape::binary_shape(self.shape(lhs), self.shape(rhs)).expect("add shape inference");
191        self.binary(BinaryOp::Add, lhs, rhs, s)
192    }
193
194    fn sub(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
195        let s = shape::binary_shape(self.shape(lhs), self.shape(rhs)).expect("sub shape inference");
196        self.binary(BinaryOp::Sub, lhs, rhs, s)
197    }
198
199    fn mul(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
200        let s = shape::binary_shape(self.shape(lhs), self.shape(rhs)).expect("mul shape inference");
201        self.binary(BinaryOp::Mul, lhs, rhs, s)
202    }
203
204    fn div(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
205        let s = shape::binary_shape(self.shape(lhs), self.shape(rhs)).expect("div shape inference");
206        self.binary(BinaryOp::Div, lhs, rhs, s)
207    }
208
209    fn gelu(&mut self, x: NodeId) -> NodeId {
210        let s = shape::unary_shape(self.shape(x));
211        self.activation(Activation::Gelu, x, s)
212    }
213
214    fn gelu_approx(&mut self, x: NodeId) -> NodeId {
215        let s = shape::unary_shape(self.shape(x));
216        self.activation(Activation::GeluApprox, x, s)
217    }
218
219    fn silu(&mut self, x: NodeId) -> NodeId {
220        let s = shape::unary_shape(self.shape(x));
221        self.activation(Activation::Silu, x, s)
222    }
223
224    fn relu(&mut self, x: NodeId) -> NodeId {
225        let s = shape::unary_shape(self.shape(x));
226        self.activation(Activation::Relu, x, s)
227    }
228
229    fn exp(&mut self, x: NodeId) -> NodeId {
230        let s = shape::unary_shape(self.shape(x));
231        self.activation(Activation::Exp, x, s)
232    }
233
234    fn sqrt(&mut self, x: NodeId) -> NodeId {
235        let s = shape::unary_shape(self.shape(x));
236        self.activation(Activation::Sqrt, x, s)
237    }
238
239    fn neg(&mut self, x: NodeId) -> NodeId {
240        let s = shape::unary_shape(self.shape(x));
241        self.activation(Activation::Neg, x, s)
242    }
243
244    fn tanh(&mut self, x: NodeId) -> NodeId {
245        let s = shape::unary_shape(self.shape(x));
246        self.activation(Activation::Tanh, x, s)
247    }
248
249    fn ln(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId {
250        let s = shape::unary_shape(self.shape(x));
251        self.layer_norm(x, gamma, beta, -1, eps, s)
252    }
253
254    fn layer_norm2d(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId {
255        Graph::layer_norm2d(self, x, gamma, beta, eps)
256    }
257
258    fn group_norm(
259        &mut self,
260        x: NodeId,
261        gamma: NodeId,
262        beta: NodeId,
263        num_groups: usize,
264        eps: f32,
265    ) -> NodeId {
266        Graph::group_norm(self, x, gamma, beta, num_groups, eps)
267    }
268
269    fn conv2d(
270        &mut self,
271        input: NodeId,
272        weight: NodeId,
273        kernel_size: [usize; 2],
274        stride: [usize; 2],
275        padding: [usize; 2],
276        dilation: [usize; 2],
277        groups: usize,
278    ) -> NodeId {
279        Graph::conv2d(
280            self,
281            input,
282            weight,
283            kernel_size,
284            stride,
285            padding,
286            dilation,
287            groups,
288        )
289    }
290
291    fn conv_transpose2d(
292        &mut self,
293        input: NodeId,
294        weight: NodeId,
295        kernel_size: [usize; 2],
296        stride: [usize; 2],
297        padding: [usize; 2],
298        dilation: [usize; 2],
299        output_padding: [usize; 2],
300        groups: usize,
301    ) -> NodeId {
302        Graph::conv_transpose2d(
303            self,
304            input,
305            weight,
306            kernel_size,
307            stride,
308            padding,
309            dilation,
310            output_padding,
311            groups,
312        )
313    }
314
315    fn rms_norm(&mut self, x: NodeId, gamma: NodeId, beta: NodeId, eps: f32) -> NodeId {
316        let s = shape::unary_shape(self.shape(x));
317        self.add_node(Op::RmsNorm { axis: -1, eps }, vec![x, gamma, beta], s)
318    }
319
320    fn sum(&mut self, x: NodeId, axes: Vec<usize>, keep_dim: bool) -> NodeId {
321        let s =
322            shape::reduce_shape(self.shape(x), &axes, keep_dim).expect("reduce shape inference");
323        self.reduce(x, ReduceOp::Sum, axes, keep_dim, s)
324    }
325
326    fn mean(&mut self, x: NodeId, axes: Vec<usize>, keep_dim: bool) -> NodeId {
327        let s =
328            shape::reduce_shape(self.shape(x), &axes, keep_dim).expect("reduce shape inference");
329        self.reduce(x, ReduceOp::Mean, axes, keep_dim, s)
330    }
331
332    fn sm(&mut self, x: NodeId, axis: i32) -> NodeId {
333        let s = shape::softmax_shape(self.shape(x));
334        self.softmax(x, axis, s)
335    }
336
337    fn reshape_(&mut self, x: NodeId, new_shape: Vec<i64>) -> NodeId {
338        let s = shape::reshape_shape(self.shape(x), &new_shape).expect("reshape shape inference");
339        self.reshape(x, new_shape, s)
340    }
341
342    fn transpose_(&mut self, x: NodeId, perm: Vec<usize>) -> NodeId {
343        let s = shape::transpose_shape(self.shape(x), &perm).expect("transpose shape inference");
344        self.add_node(Op::Transpose { perm }, vec![x], s)
345    }
346
347    fn narrow_(&mut self, x: NodeId, axis: usize, start: usize, len: usize) -> NodeId {
348        let s = shape::narrow_shape(self.shape(x), axis, len).expect("narrow shape inference");
349        self.add_node(Op::Narrow { axis, start, len }, vec![x], s)
350    }
351
352    fn concat_(&mut self, inputs: Vec<NodeId>, axis: usize) -> NodeId {
353        let shapes: Vec<&Shape> = inputs.iter().map(|&id| self.shape(id)).collect();
354        let s = shape::concat_shape(&shapes, axis).expect("concat shape inference");
355        self.concat(inputs, axis, s)
356    }
357
358    fn gather_(&mut self, table: NodeId, indices: NodeId, axis: usize) -> NodeId {
359        let s = shape::gather_shape(self.shape(table), self.shape(indices), axis)
360            .expect("gather shape inference");
361        self.gather(table, indices, axis, s)
362    }
363
364    fn eq(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
365        let s = shape::compare_shape(self.shape(lhs), self.shape(rhs))
366            .expect("compare shape inference");
367        self.add_node(Op::Compare(CmpOp::Eq), vec![lhs, rhs], s)
368    }
369
370    fn lt(&mut self, lhs: NodeId, rhs: NodeId) -> NodeId {
371        let s = shape::compare_shape(self.shape(lhs), self.shape(rhs))
372            .expect("compare shape inference");
373        self.add_node(Op::Compare(CmpOp::Lt), vec![lhs, rhs], s)
374    }
375
376    fn attention_(
377        &mut self,
378        q: NodeId,
379        k: NodeId,
380        v: NodeId,
381        mask: NodeId,
382        num_heads: usize,
383        head_dim: usize,
384    ) -> NodeId {
385        let s = shape::attention_shape(self.shape(q));
386        self.attention(q, k, v, mask, num_heads, head_dim, s)
387    }
388
389    fn rope(&mut self, x: NodeId, cos: NodeId, sin: NodeId, head_dim: usize) -> NodeId {
390        self.rope_n(x, cos, sin, head_dim, head_dim)
391    }
392
393    fn rope_n(
394        &mut self,
395        x: NodeId,
396        cos: NodeId,
397        sin: NodeId,
398        head_dim: usize,
399        n_rot: usize,
400    ) -> NodeId {
401        self.rope_n_styled(x, cos, sin, head_dim, n_rot, crate::op::RopeStyle::NeoX)
402    }
403
404    fn rope_styled(
405        &mut self,
406        x: NodeId,
407        cos: NodeId,
408        sin: NodeId,
409        head_dim: usize,
410        style: crate::op::RopeStyle,
411    ) -> NodeId {
412        self.rope_n_styled(x, cos, sin, head_dim, head_dim, style)
413    }
414
415    fn rope_n_styled(
416        &mut self,
417        x: NodeId,
418        cos: NodeId,
419        sin: NodeId,
420        head_dim: usize,
421        n_rot: usize,
422        style: crate::op::RopeStyle,
423    ) -> NodeId {
424        assert!(
425            n_rot <= head_dim && n_rot.is_multiple_of(2),
426            "rope_n: n_rot={n_rot} must be even and <= head_dim={head_dim}"
427        );
428        let s = shape::unary_shape(self.shape(x));
429        self.add_node(
430            Op::Rope {
431                head_dim,
432                n_rot,
433                style,
434            },
435            vec![x, cos, sin],
436            s,
437        )
438    }
439
440    fn cast(&mut self, x: NodeId, to: DType) -> NodeId {
441        let s = shape::cast_shape(self.shape(x), to);
442        self.add_node(Op::Cast { to }, vec![x], s)
443    }
444
445    fn try_constant(&mut self, value: f64, dtype: DType) -> Result<NodeId, String> {
446        if matches!(dtype, DType::F16 | DType::BF16) {
447            let f32_id = self.try_constant(value, DType::F32)?;
448            return Ok(self.cast(f32_id, dtype));
449        }
450        let data = scalar_constant_bytes(value, dtype)?;
451        Ok(self.add_node(Op::Constant { data }, vec![], Shape::scalar(dtype)))
452    }
453
454    fn constant(&mut self, value: f64, dtype: DType) -> NodeId {
455        self.try_constant(value, dtype)
456            .expect("scalar constant encoding")
457    }
458
459    fn stop_gradient(&mut self, x: NodeId) -> NodeId {
460        let s = shape::unary_shape(self.shape(x));
461        self.add_node(Op::StopGradient, vec![x], s)
462    }
463}
464
465#[cfg(test)]
466mod tests {
467    use super::*;
468
469    #[test]
470    fn inferred_conv2d_and_conv_transpose2d() {
471        let mut g = Graph::new("conv");
472        let f = DType::F32;
473        let x = g.input("x", Shape::new(&[1, 4, 8, 8], f));
474        let w = g.param("w", Shape::new(&[8, 2, 3, 3], f));
475        let y = g.conv2d(x, w, [3, 3], [1, 1], [1, 1], [1, 1], 2);
476        assert_eq!(g.shape(y), &Shape::new(&[1, 8, 8, 8], f));
477
478        let wt = g.param("wt", Shape::new(&[4, 8, 2, 2], f));
479        let z = g.conv_transpose2d(x, wt, [2, 2], [2, 2], [0, 0], [1, 1], [0, 0], 1);
480        assert_eq!(g.shape(z), &Shape::new(&[1, 8, 16, 16], f));
481    }
482
483    #[test]
484    fn inferred_layer_norm2d() {
485        let mut g = Graph::new("ln2d");
486        let f = DType::F32;
487        let x = g.input("x", Shape::new(&[1, 4, 8, 8], f));
488        let gamma = g.param("g", Shape::new(&[4], f));
489        let beta = g.param("b", Shape::new(&[4], f));
490        let y = g.layer_norm2d(x, gamma, beta, 1e-6);
491        assert_eq!(g.shape(y), &Shape::new(&[1, 4, 8, 8], f));
492    }
493
494    #[test]
495    fn inferred_matmul_bias_gelu() {
496        let mut g = Graph::new("test");
497        let x = g.input("x", Shape::new(&[4, 15, 384], DType::F32));
498        let w = g.param("w", Shape::new(&[384, 1536], DType::F32));
499        let b = g.param("b", Shape::new(&[1536], DType::F32));
500
501        // No explicit shapes needed!
502        let mm = g.mm(x, w);
503        let add = g.add(mm, b);
504        let out = g.gelu(add);
505        g.set_outputs(vec![out]);
506
507        assert_eq!(g.shape(mm), &Shape::new(&[4, 15, 1536], DType::F32));
508        assert_eq!(g.shape(add), &Shape::new(&[4, 15, 1536], DType::F32));
509        assert_eq!(g.shape(out), &Shape::new(&[4, 15, 1536], DType::F32));
510    }
511
512    #[test]
513    fn inferred_bert_ffn() {
514        let mut g = Graph::new("bert_ffn");
515        let f = DType::F32;
516        let h = 384;
517        let int = 1536;
518
519        let x = g.input("x", Shape::new(&[4, 15, h], f));
520        let int_w = g.param("int.w", Shape::new(&[h, int], f));
521        let int_b = g.param("int.b", Shape::new(&[int], f));
522        let out_w = g.param("out.w", Shape::new(&[int, h], f));
523        let out_b = g.param("out.b", Shape::new(&[h], f));
524        let gamma = g.param("g", Shape::new(&[h], f));
525        let beta = g.param("b", Shape::new(&[h], f));
526
527        let mm1 = g.mm(x, int_w);
528        let a1 = g.add(mm1, int_b);
529        let ffn = g.gelu(a1);
530        let mm2 = g.mm(ffn, out_w);
531        let out = g.add(mm2, out_b);
532        let res = g.add(out, x);
533        let normed = g.ln(res, gamma, beta, 1e-12);
534        g.set_outputs(vec![normed]);
535
536        assert_eq!(g.shape(ffn), &Shape::new(&[4, 15, int], f));
537        assert_eq!(g.shape(out), &Shape::new(&[4, 15, h], f));
538        assert_eq!(g.shape(normed), &Shape::new(&[4, 15, h], f));
539    }
540
541    #[test]
542    fn inferred_gather_reshape() {
543        let mut g = Graph::new("test");
544        let table = g.param("emb", Shape::new(&[30522, 384], DType::F32));
545        let ids = g.input("ids", Shape::new(&[4, 15], DType::I64));
546
547        let gathered = g.gather_(table, ids, 0);
548        assert_eq!(g.shape(gathered), &Shape::new(&[4, 15, 384], DType::F32));
549
550        let reshaped = g.reshape_(gathered, vec![60, 384]);
551        assert_eq!(g.shape(reshaped), &Shape::new(&[60, 384], DType::F32));
552
553        let transposed = g.transpose_(reshaped, vec![1, 0]);
554        assert_eq!(g.shape(transposed), &Shape::new(&[384, 60], DType::F32));
555    }
556
557    #[test]
558    fn inferred_constant_broadcasts() {
559        let mut g = Graph::new("const");
560        let x = g.input("x", Shape::new(&[2, 3], DType::F32));
561        let c = g.constant(2.0, DType::F32);
562        assert_eq!(g.shape(c), &Shape::scalar(DType::F32));
563        let y = g.mul(x, c);
564        assert_eq!(g.shape(y), &Shape::new(&[2, 3], DType::F32));
565    }
566
567    #[test]
568    fn inferred_constant_f16_via_cast() {
569        let mut g = Graph::new("const_f16");
570        let c = g.constant(1.5, DType::F16);
571        assert_eq!(g.shape(c), &Shape::scalar(DType::F16));
572        let x = g.input("x", Shape::new(&[2], DType::F16));
573        let y = g.add(x, c);
574        assert_eq!(g.shape(y), &Shape::new(&[2], DType::F16));
575    }
576
577    #[test]
578    fn inferred_constant_arithmetic_chain() {
579        let mut g = Graph::new("const_chain");
580        let x = g.input("x", Shape::new(&[4], DType::F32));
581        let one = g.constant(1.0, DType::F32);
582        let two = g.constant(2.0, DType::F32);
583        let sum = g.add(x, one);
584        let y = g.div(sum, two);
585        assert_eq!(g.shape(y), &Shape::new(&[4], DType::F32));
586        g.set_outputs(vec![y]);
587    }
588
589    #[test]
590    fn try_constant_rejects_out_of_range() {
591        let mut g = Graph::new("try_const");
592        let err = g.try_constant(128.0, DType::I8).unwrap_err();
593        assert!(err.contains("out of range"));
594    }
595
596    #[test]
597    fn try_constant_f16_via_cast() {
598        let mut g = Graph::new("try_const_f16");
599        let c = g.try_constant(1.5, DType::F16).unwrap();
600        assert_eq!(g.shape(c), &Shape::scalar(DType::F16));
601    }
602
603    #[test]
604    fn inferred_reduce_softmax() {
605        let mut g = Graph::new("test");
606        let x = g.input("x", Shape::new(&[4, 15, 384], DType::F32));
607
608        let s = g.sm(x, -1);
609        assert_eq!(g.shape(s), &Shape::new(&[4, 15, 384], DType::F32));
610
611        let m = g.mean(x, vec![2], false);
612        assert_eq!(g.shape(m), &Shape::new(&[4, 15], DType::F32));
613
614        let mk = g.mean(x, vec![2], true);
615        assert_eq!(g.shape(mk), &Shape::new(&[4, 15, 1], DType::F32));
616    }
617}
rlx_ir/infer.rs

rlx_ir/
infer.rs
