lift_tensor/
shape.rs

1use crate::ops::TensorOp;
2use lift_core::types::{Dimension, TensorTypeInfo};
3
4#[derive(Debug)]
5pub struct ShapeInference;
6
7impl ShapeInference {
8    pub fn infer_output_shape(
9        op: &TensorOp,
10        inputs: &[&TensorTypeInfo],
11    ) -> Result<Vec<TensorTypeInfo>, String> {
12        match op {
13            // ── Binary element-wise (broadcast) ──
14            TensorOp::Add | TensorOp::Sub | TensorOp::Mul | TensorOp::Div => {
15                if inputs.len() != 2 {
16                    return Err(format!("{} requires 2 inputs", op.name()));
17                }
18                let result = broadcast_shapes(&inputs[0].shape, &inputs[1].shape)?;
19                Ok(vec![TensorTypeInfo {
20                    shape: result,
21                    dtype: inputs[0].dtype,
22                    layout: inputs[0].layout,
23                }])
24            }
25
26            // ── Unary shape-preserving ──
27            TensorOp::Neg
28            | TensorOp::ReLU
29            | TensorOp::GeLU
30            | TensorOp::SiLU
31            | TensorOp::Sigmoid
32            | TensorOp::Tanh
33            | TensorOp::LeakyReLU
34            | TensorOp::ELU
35            | TensorOp::Mish
36            | TensorOp::HardSwish
37            | TensorOp::HardSigmoid
38            | TensorOp::Softmax
39            | TensorOp::Cumsum
40            | TensorOp::Quantize
41            | TensorOp::Dequantize
42            | TensorOp::QuantizeInt4
43            | TensorOp::DequantizeInt4
44            | TensorOp::QuantizeFp8
45            | TensorOp::DequantizeFp8
46            | TensorOp::Checkpoint
47            | TensorOp::Offload
48            | TensorOp::GradReLU
49            | TensorOp::GradGeLU
50            | TensorOp::GradSoftmax => {
51                if inputs.is_empty() {
52                    return Err(format!("{} requires at least 1 input", op.name()));
53                }
54                Ok(vec![inputs[0].clone()])
55            }
56
57            // ── Normalisation (shape-preserving) ──
58            TensorOp::LayerNorm
59            | TensorOp::RMSNorm
60            | TensorOp::BatchNorm
61            | TensorOp::GroupNorm
62            | TensorOp::InstanceNorm
63            | TensorOp::GradLayerNorm => {
64                if inputs.is_empty() {
65                    return Err(format!("{} requires at least 1 input", op.name()));
66                }
67                Ok(vec![inputs[0].clone()])
68            }
69
70            // ── MatMul ──
71            TensorOp::MatMul | TensorOp::SparseMatMul => {
72                if inputs.len() != 2 {
73                    return Err("matmul requires 2 inputs".into());
74                }
75                let a = &inputs[0].shape;
76                let b = &inputs[1].shape;
77                if a.len() < 2 || b.len() < 2 {
78                    return Err("matmul inputs must be at least 2D".into());
79                }
80                let m = a[a.len() - 2].clone();
81                let n = b[b.len() - 1].clone();
82
83                let k_a = &a[a.len() - 1];
84                let k_b = &b[b.len() - 2];
85                if let (Some(ka), Some(kb)) = (k_a.static_value(), k_b.static_value()) {
86                    if ka != kb {
87                        return Err(format!("matmul inner dimension mismatch: {} vs {}", ka, kb));
88                    }
89                }
90
91                let mut result_shape = Vec::new();
92                let batch_a = &a[..a.len() - 2];
93                let batch_b = &b[..b.len() - 2];
94                let batch = broadcast_shapes(batch_a, batch_b)?;
95                result_shape.extend(batch);
96                result_shape.push(m);
97                result_shape.push(n);
98
99                Ok(vec![TensorTypeInfo {
100                    shape: result_shape,
101                    dtype: inputs[0].dtype,
102                    layout: inputs[0].layout,
103                }])
104            }
105
106            // ── Linear ──
107            TensorOp::Linear => {
108                if inputs.len() < 2 {
109                    return Err("linear requires at least 2 inputs (x, W)".into());
110                }
111                let x = &inputs[0].shape;
112                let w = &inputs[1].shape;
113                if x.is_empty() || w.len() != 2 {
114                    return Err("linear: x must be at least 1D, W must be 2D".into());
115                }
116                let mut result_shape = x[..x.len() - 1].to_vec();
117                result_shape.push(w[1].clone());
118
119                Ok(vec![TensorTypeInfo {
120                    shape: result_shape,
121                    dtype: inputs[0].dtype,
122                    layout: inputs[0].layout,
123                }])
124            }
125
126            // ── Conv2D ──
127            TensorOp::Conv2D | TensorOp::DepthwiseConv2D | TensorOp::DilatedConv2D => {
128                if inputs.len() < 2 {
129                    return Err("conv2d requires at least 2 inputs (input, kernel)".into());
130                }
131                let input = &inputs[0].shape;
132                let kernel = &inputs[1].shape;
133                if input.len() != 4 || kernel.len() != 4 {
134                    return Err("conv2d: input and kernel must be 4D (NCHW)".into());
135                }
136
137                let n = input[0].clone();
138                let cout = kernel[0].clone();
139                let h_out = match (&input[2], &kernel[2]) {
140                    (Dimension::Constant(ih), Dimension::Constant(kh)) => {
141                        Dimension::Constant(ih - kh + 1)
142                    }
143                    _ => Dimension::Symbolic("H_out".into()),
144                };
145                let w_out = match (&input[3], &kernel[3]) {
146                    (Dimension::Constant(iw), Dimension::Constant(kw)) => {
147                        Dimension::Constant(iw - kw + 1)
148                    }
149                    _ => Dimension::Symbolic("W_out".into()),
150                };
151
152                Ok(vec![TensorTypeInfo {
153                    shape: vec![n, cout, h_out, w_out],
154                    dtype: inputs[0].dtype,
155                    layout: inputs[0].layout,
156                }])
157            }
158
159            // ── Conv1D ──
160            TensorOp::Conv1D => {
161                if inputs.len() < 2 {
162                    return Err("conv1d requires at least 2 inputs".into());
163                }
164                let input = &inputs[0].shape;
165                let kernel = &inputs[1].shape;
166                if input.len() != 3 || kernel.len() != 3 {
167                    return Err("conv1d: input [N,C,L] and kernel [Cout,Cin,K]".into());
168                }
169                let n = input[0].clone();
170                let cout = kernel[0].clone();
171                let l_out = match (&input[2], &kernel[2]) {
172                    (Dimension::Constant(il), Dimension::Constant(kl)) => {
173                        Dimension::Constant(il - kl + 1)
174                    }
175                    _ => Dimension::Symbolic("L_out".into()),
176                };
177                Ok(vec![TensorTypeInfo {
178                    shape: vec![n, cout, l_out],
179                    dtype: inputs[0].dtype,
180                    layout: inputs[0].layout,
181                }])
182            }
183
184            // ── Conv3D ──
185            TensorOp::Conv3D => {
186                if inputs.len() < 2 {
187                    return Err("conv3d requires at least 2 inputs".into());
188                }
189                let input = &inputs[0].shape;
190                let kernel = &inputs[1].shape;
191                if input.len() != 5 || kernel.len() != 5 {
192                    return Err("conv3d: input [N,C,D,H,W] and kernel [Cout,Cin,Kd,Kh,Kw]".into());
193                }
194                let n = input[0].clone();
195                let cout = kernel[0].clone();
196                let dims: Vec<Dimension> = (2..5)
197                    .map(|i| match (&input[i], &kernel[i]) {
198                        (Dimension::Constant(iv), Dimension::Constant(kv)) => {
199                            Dimension::Constant(iv - kv + 1)
200                        }
201                        _ => Dimension::Symbolic(format!("dim{}_out", i)),
202                    })
203                    .collect();
204                Ok(vec![TensorTypeInfo {
205                    shape: vec![n, cout, dims[0].clone(), dims[1].clone(), dims[2].clone()],
206                    dtype: inputs[0].dtype,
207                    layout: inputs[0].layout,
208                }])
209            }
210
211            // ── Pooling ──
212            TensorOp::MaxPool2D | TensorOp::AvgPool2D => {
213                if inputs.is_empty() {
214                    return Err(format!("{} requires at least 1 input", op.name()));
215                }
216                // Simplified: returns same shape (caller should use attrs for kernel/stride)
217                Ok(vec![inputs[0].clone()])
218            }
219
220            TensorOp::AdaptiveAvgPool2D => {
221                if inputs.is_empty() {
222                    return Err("adaptive_avgpool2d requires 1 input".into());
223                }
224                Ok(vec![inputs[0].clone()])
225            }
226
227            TensorOp::GlobalAvgPool => {
228                if inputs.is_empty() {
229                    return Err("global_avgpool requires 1 input".into());
230                }
231                let shape = &inputs[0].shape;
232                if shape.len() < 3 {
233                    return Err("global_avgpool: input must be at least 3D [N,C,...]".into());
234                }
235                // [N, C, ...] -> [N, C, 1, 1, ...]
236                let mut out = vec![shape[0].clone(), shape[1].clone()];
237                for _ in 2..shape.len() {
238                    out.push(Dimension::Constant(1));
239                }
240                Ok(vec![TensorTypeInfo {
241                    shape: out,
242                    dtype: inputs[0].dtype,
243                    layout: inputs[0].layout,
244                }])
245            }
246
247            // ── Attention variants ──
248            TensorOp::Attention
249            | TensorOp::MultiHeadAttention
250            | TensorOp::MultiQueryAttention
251            | TensorOp::GroupedQueryAttention
252            | TensorOp::FlashAttention
253            | TensorOp::SlidingWindowAttention
254            | TensorOp::CrossAttention
255            | TensorOp::PagedAttention
256            | TensorOp::GradAttention => {
257                if inputs.len() < 3 {
258                    return Err("attention requires at least 3 inputs (Q, K, V)".into());
259                }
260                Ok(vec![inputs[0].clone()])
261            }
262
263            // ── Recurrent ──
264            TensorOp::LSTMCell => {
265                if inputs.len() < 2 {
266                    return Err("lstm_cell requires input and hidden state".into());
267                }
268                // Returns (h_new, c_new) with same shape as hidden
269                Ok(vec![inputs[1].clone(), inputs[1].clone()])
270            }
271
272            TensorOp::GRUCell | TensorOp::RNNCell => {
273                if inputs.len() < 2 {
274                    return Err(format!("{} requires input and hidden state", op.name()));
275                }
276                Ok(vec![inputs[1].clone()])
277            }
278
279            // ── Shape / zero-flop ops ──
280            TensorOp::Reshape
281            | TensorOp::Transpose
282            | TensorOp::Squeeze
283            | TensorOp::Unsqueeze
284            | TensorOp::Permute
285            | TensorOp::Expand
286            | TensorOp::Slice
287            | TensorOp::Pad
288            | TensorOp::Tile => {
289                // These need target shape from attributes; passthrough for now
290                if inputs.is_empty() {
291                    return Err(format!("{} requires at least 1 input", op.name()));
292                }
293                Ok(vec![inputs[0].clone()])
294            }
295
296            // ── Concat ──
297            TensorOp::Concat => {
298                if inputs.is_empty() {
299                    return Err("concat requires at least 1 input".into());
300                }
301                Ok(vec![inputs[0].clone()])
302            }
303
304            // ── TopK / Sort ──
305            TensorOp::TopK | TensorOp::Sort => {
306                if inputs.is_empty() {
307                    return Err(format!("{} requires 1 input", op.name()));
308                }
309                Ok(vec![inputs[0].clone()])
310            }
311
312            // ── FFT / IFFT ──
313            TensorOp::FFT | TensorOp::IFFT => {
314                if inputs.is_empty() {
315                    return Err(format!("{} requires 1 input", op.name()));
316                }
317                Ok(vec![inputs[0].clone()])
318            }
319
320            // ── SVD: returns U, S, V ──
321            TensorOp::SVD => {
322                if inputs.is_empty() {
323                    return Err("svd requires 1 input".into());
324                }
325                Ok(vec![inputs[0].clone()])
326            }
327
328            // ── Where: condition, x, y -> x ──
329            TensorOp::Where | TensorOp::Clamp => {
330                if inputs.len() < 2 {
331                    return Err(format!("{} requires at least 2 inputs", op.name()));
332                }
333                Ok(vec![inputs[0].clone()])
334            }
335
336            _ => {
337                // For ops not yet handled, passthrough first input or empty
338                if !inputs.is_empty() {
339                    Ok(vec![inputs[0].clone()])
340                } else {
341                    Ok(Vec::new())
342                }
343            }
344        }
345    }
346
347    pub fn compute_flops(op: &TensorOp, inputs: &[&TensorTypeInfo]) -> Option<u64> {
348        match op {
349            TensorOp::MatMul | TensorOp::SparseMatMul => {
350                if inputs.len() != 2 {
351                    return None;
352                }
353                let a = &inputs[0].shape;
354                let b = &inputs[1].shape;
355                let m = a.get(a.len().checked_sub(2)?)?.static_value()? as u64;
356                let k = a.last()?.static_value()? as u64;
357                let n = b.last()?.static_value()? as u64;
358                let batch: u64 = a[..a.len() - 2]
359                    .iter()
360                    .filter_map(|d| d.static_value())
361                    .map(|v| v as u64)
362                    .product::<u64>()
363                    .max(1);
364                Some(2 * batch * m * n * k)
365            }
366
367            TensorOp::Add | TensorOp::Sub | TensorOp::Mul | TensorOp::Div => {
368                if inputs.is_empty() {
369                    return None;
370                }
371                Some(element_count(&inputs[0].shape)? as u64)
372            }
373
374            TensorOp::ReLU
375            | TensorOp::Sigmoid
376            | TensorOp::Tanh
377            | TensorOp::LeakyReLU
378            | TensorOp::ELU
379            | TensorOp::HardSigmoid => {
380                if inputs.is_empty() {
381                    return None;
382                }
383                Some(element_count(&inputs[0].shape)? as u64)
384            }
385
386            TensorOp::GeLU | TensorOp::SiLU | TensorOp::Mish | TensorOp::HardSwish => {
387                if inputs.is_empty() {
388                    return None;
389                }
390                let n = element_count(&inputs[0].shape)? as u64;
391                Some(8 * n)
392            }
393
394            TensorOp::Softmax => {
395                if inputs.is_empty() {
396                    return None;
397                }
398                let n = element_count(&inputs[0].shape)? as u64;
399                Some(5 * n)
400            }
401
402            TensorOp::LayerNorm
403            | TensorOp::RMSNorm
404            | TensorOp::GroupNorm
405            | TensorOp::InstanceNorm => {
406                if inputs.is_empty() {
407                    return None;
408                }
409                let n = element_count(&inputs[0].shape)? as u64;
410                Some(7 * n)
411            }
412
413            TensorOp::BatchNorm => {
414                if inputs.is_empty() {
415                    return None;
416                }
417                let n = element_count(&inputs[0].shape)? as u64;
418                Some(5 * n)
419            }
420
421            TensorOp::Linear => {
422                if inputs.len() < 2 {
423                    return None;
424                }
425                let x = &inputs[0].shape;
426                let w = &inputs[1].shape;
427                let m: u64 = x[..x.len() - 1]
428                    .iter()
429                    .filter_map(|d| d.static_value())
430                    .map(|v| v as u64)
431                    .product::<u64>()
432                    .max(1);
433                let k = x.last()?.static_value()? as u64;
434                let n = w.last()?.static_value()? as u64;
435                Some(2 * m * n * k + n)
436            }
437
438            TensorOp::Conv2D | TensorOp::DepthwiseConv2D | TensorOp::DilatedConv2D => {
439                if inputs.len() < 2 {
440                    return None;
441                }
442                let kernel = &inputs[1].shape;
443                let cout = kernel[0].static_value()? as u64;
444                let cin = kernel[1].static_value()? as u64;
445                let kh = kernel[2].static_value()? as u64;
446                let kw = kernel[3].static_value()? as u64;
447                let input = &inputs[0].shape;
448                let n = input[0].static_value()? as u64;
449                let ih = input[2].static_value()? as u64;
450                let iw = input[3].static_value()? as u64;
451                let oh = ih.saturating_sub(kh) + 1;
452                let ow = iw.saturating_sub(kw) + 1;
453                Some(2 * n * cout * cin * kh * kw * oh * ow)
454            }
455
456            TensorOp::Conv1D => {
457                if inputs.len() < 2 {
458                    return None;
459                }
460                let kernel = &inputs[1].shape;
461                let cout = kernel[0].static_value()? as u64;
462                let cin = kernel[1].static_value()? as u64;
463                let k = kernel[2].static_value()? as u64;
464                let input = &inputs[0].shape;
465                let n = input[0].static_value()? as u64;
466                let il = input[2].static_value()? as u64;
467                let ol = il.saturating_sub(k) + 1;
468                Some(2 * n * cout * cin * k * ol)
469            }
470
471            TensorOp::Conv3D => {
472                if inputs.len() < 2 {
473                    return None;
474                }
475                let kernel = &inputs[1].shape;
476                let cout = kernel.first()?.static_value()? as u64;
477                let cin = kernel.get(1)?.static_value()? as u64;
478                let kd = kernel.get(2)?.static_value()? as u64;
479                let kh = kernel.get(3)?.static_value()? as u64;
480                let kw = kernel.get(4)?.static_value()? as u64;
481                let input = &inputs[0].shape;
482                let n = input.first()?.static_value()? as u64;
483                let id = input.get(2)?.static_value()? as u64;
484                let ih = input.get(3)?.static_value()? as u64;
485                let iw = input.get(4)?.static_value()? as u64;
486                let od = id.saturating_sub(kd) + 1;
487                let oh = ih.saturating_sub(kh) + 1;
488                let ow = iw.saturating_sub(kw) + 1;
489                Some(2 * n * cout * cin * kd * kh * kw * od * oh * ow)
490            }
491
492            // Attention variants: 2*B*H*(S^2*D + S*D^2)
493            TensorOp::Attention
494            | TensorOp::MultiHeadAttention
495            | TensorOp::MultiQueryAttention
496            | TensorOp::GroupedQueryAttention
497            | TensorOp::FlashAttention
498            | TensorOp::SlidingWindowAttention
499            | TensorOp::CrossAttention => {
500                if inputs.is_empty() {
501                    return None;
502                }
503                let shape = &inputs[0].shape;
504                if shape.len() < 3 {
505                    return None;
506                }
507                let b = shape[0].static_value().unwrap_or(1) as u64;
508                let s = shape[shape.len() - 2].static_value()? as u64;
509                let d = shape.last()?.static_value()? as u64;
510                let h = if shape.len() >= 4 {
511                    shape[1].static_value().unwrap_or(1) as u64
512                } else {
513                    1
514                };
515                Some(4 * b * h * s * s * d)
516            }
517
518            // Recurrent
519            TensorOp::LSTMCell => {
520                // 4 * (input_size + hidden_size) * hidden_size * 2
521                if inputs.len() < 2 {
522                    return None;
523                }
524                let input_size = inputs[0].shape.last()?.static_value()? as u64;
525                let hidden_size = inputs[1].shape.last()?.static_value()? as u64;
526                Some(8 * (input_size + hidden_size) * hidden_size)
527            }
528
529            TensorOp::GRUCell => {
530                if inputs.len() < 2 {
531                    return None;
532                }
533                let input_size = inputs[0].shape.last()?.static_value()? as u64;
534                let hidden_size = inputs[1].shape.last()?.static_value()? as u64;
535                Some(6 * (input_size + hidden_size) * hidden_size)
536            }
537
538            TensorOp::RNNCell => {
539                if inputs.len() < 2 {
540                    return None;
541                }
542                let input_size = inputs[0].shape.last()?.static_value()? as u64;
543                let hidden_size = inputs[1].shape.last()?.static_value()? as u64;
544                Some(2 * (input_size + hidden_size) * hidden_size)
545            }
546
547            // FFT: 5*N*log2(N)
548            TensorOp::FFT | TensorOp::IFFT => {
549                if inputs.is_empty() {
550                    return None;
551                }
552                let n = element_count(&inputs[0].shape)? as u64;
553                if n == 0 {
554                    return Some(0);
555                }
556                let log2n = (n as f64).log2().ceil() as u64;
557                Some(5 * n * log2n)
558            }
559
560            // Pooling
561            TensorOp::MaxPool2D
562            | TensorOp::AvgPool2D
563            | TensorOp::AdaptiveAvgPool2D
564            | TensorOp::GlobalAvgPool => {
565                if inputs.is_empty() {
566                    return None;
567                }
568                Some(element_count(&inputs[0].shape)? as u64)
569            }
570
571            // Zero-flop ops
572            _ if op.is_zero_flop() => Some(0),
573
574            _ => None,
575        }
576    }
577
578    pub fn compute_memory_bytes(op: &TensorOp, inputs: &[&TensorTypeInfo]) -> Option<u64> {
579        match op {
580            TensorOp::MatMul | TensorOp::SparseMatMul => {
581                if inputs.len() != 2 {
582                    return None;
583                }
584                let a_bytes = tensor_bytes(inputs[0])? as u64;
585                let b_bytes = tensor_bytes(inputs[1])? as u64;
586                let out_shape = Self::infer_output_shape(op, inputs).ok()?;
587                let out_bytes = if let Some(out) = out_shape.first() {
588                    tensor_info_bytes(out)? as u64
589                } else {
590                    0
591                };
592                Some(a_bytes + b_bytes + out_bytes)
593            }
594            _ => {
595                let total: u64 = inputs
596                    .iter()
597                    .filter_map(|i| tensor_bytes(i).map(|b| b as u64))
598                    .sum();
599                Some(total)
600            }
601        }
602    }
603}
604
605fn broadcast_shapes(a: &[Dimension], b: &[Dimension]) -> Result<Vec<Dimension>, String> {
606    let max_rank = a.len().max(b.len());
607    let mut result = Vec::with_capacity(max_rank);
608
609    for i in 0..max_rank {
610        let da = if i < a.len() {
611            Some(&a[a.len() - 1 - i])
612        } else {
613            None
614        };
615        let db = if i < b.len() {
616            Some(&b[b.len() - 1 - i])
617        } else {
618            None
619        };
620
621        let dim = match (da, db) {
622            (Some(a_dim), Some(b_dim)) => match (a_dim.static_value(), b_dim.static_value()) {
623                (Some(a_val), Some(b_val)) => {
624                    if a_val == b_val {
625                        Dimension::Constant(a_val)
626                    } else if a_val == 1 {
627                        Dimension::Constant(b_val)
628                    } else if b_val == 1 {
629                        Dimension::Constant(a_val)
630                    } else {
631                        return Err(format!("Shape broadcast error: {} vs {}", a_val, b_val));
632                    }
633                }
634                _ => Dimension::Symbolic("broadcast".into()),
635            },
636            (Some(d), None) | (None, Some(d)) => d.clone(),
637            (None, None) => unreachable!(),
638        };
639        result.push(dim);
640    }
641
642    result.reverse();
643    Ok(result)
644}
645
646fn element_count(shape: &[Dimension]) -> Option<usize> {
647    let mut count = 1usize;
648    for dim in shape {
649        count = count.checked_mul(dim.static_value()?)?;
650    }
651    Some(count)
652}
653
654fn tensor_bytes(info: &TensorTypeInfo) -> Option<usize> {
655    Some(element_count(&info.shape)? * info.dtype.byte_size())
656}
657
658fn tensor_info_bytes(info: &TensorTypeInfo) -> Option<usize> {
659    tensor_bytes(info)
660}
661
662#[cfg(test)]
663mod tests {
664    use super::*;
665    use lift_core::types::{DataType, MemoryLayout};
666
667    fn make_tensor(shape: Vec<usize>, dtype: DataType) -> TensorTypeInfo {
668        TensorTypeInfo {
669            shape: shape.into_iter().map(Dimension::Constant).collect(),
670            dtype,
671            layout: MemoryLayout::Contiguous,
672        }
673    }
674
675    #[test]
676    fn test_matmul_shape() {
677        let a = make_tensor(vec![2, 3, 4], DataType::FP32);
678        let b = make_tensor(vec![2, 4, 5], DataType::FP32);
679        let result = ShapeInference::infer_output_shape(&TensorOp::MatMul, &[&a, &b]).unwrap();
680        assert_eq!(result.len(), 1);
681        let shape = &result[0].shape;
682        assert_eq!(shape.len(), 3);
683        assert_eq!(shape[0].static_value(), Some(2));
684        assert_eq!(shape[1].static_value(), Some(3));
685        assert_eq!(shape[2].static_value(), Some(5));
686    }
687
688    #[test]
689    fn test_matmul_dimension_mismatch() {
690        let a = make_tensor(vec![3, 4], DataType::FP32);
691        let b = make_tensor(vec![5, 6], DataType::FP32);
692        let result = ShapeInference::infer_output_shape(&TensorOp::MatMul, &[&a, &b]);
693        assert!(result.is_err());
694    }
695
696    #[test]
697    fn test_matmul_flops() {
698        let a = make_tensor(vec![2, 3], DataType::FP32);
699        let b = make_tensor(vec![3, 4], DataType::FP32);
700        let flops = ShapeInference::compute_flops(&TensorOp::MatMul, &[&a, &b]);
701        assert_eq!(flops, Some(2 * 2 * 4 * 3)); // 2*M*N*K
702    }
703
704    #[test]
705    fn test_relu_shape() {
706        let a = make_tensor(vec![2, 3, 4], DataType::FP32);
707        let result = ShapeInference::infer_output_shape(&TensorOp::ReLU, &[&a]).unwrap();
708        assert_eq!(result[0].shape, a.shape);
709    }
710
711    #[test]
712    fn test_linear_shape() {
713        let x = make_tensor(vec![1, 784], DataType::FP32);
714        let w = make_tensor(vec![784, 64], DataType::FP32);
715        let b = make_tensor(vec![64], DataType::FP32);
716        let result = ShapeInference::infer_output_shape(&TensorOp::Linear, &[&x, &w, &b]).unwrap();
717        assert_eq!(result[0].shape[0].static_value(), Some(1));
718        assert_eq!(result[0].shape[1].static_value(), Some(64));
719    }
720
721    #[test]
722    fn test_conv2d_shape() {
723        let input = make_tensor(vec![1, 3, 28, 28], DataType::FP32);
724        let kernel = make_tensor(vec![16, 3, 5, 5], DataType::FP32);
725        let result =
726            ShapeInference::infer_output_shape(&TensorOp::Conv2D, &[&input, &kernel]).unwrap();
727        assert_eq!(result[0].shape[0].static_value(), Some(1));
728        assert_eq!(result[0].shape[1].static_value(), Some(16));
729        assert_eq!(result[0].shape[2].static_value(), Some(24)); // 28-5+1
730        assert_eq!(result[0].shape[3].static_value(), Some(24));
731    }
732}
lift_tensor/shape.rs

lift_tensor/
shape.rs
