1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
//! Broadcasting operations for tensors.
//!
//! Implements NumPy-style broadcasting rules:
//! - Shapes are aligned from the right (trailing dimensions)
//! - Missing dimensions are treated as 1
//! - For each dimension, sizes must either match or one must be 1
//!
//! This module provides the infrastructure for automatic broadcasting
//! in binary operations, matching Tinygrad's architecture.
use snafu::ResultExt;
use super::*;
use svod_ir::shape::{align_shapes_left, broadcast_shape};
impl Tensor {
/// Broadcast two tensors to a common shape for binary operations.
///
/// This method implements automatic broadcasting similar to NumPy/PyTorch.
/// It aligns the shapes, computes the broadcast result shape, and broadcasts
/// each tensor to that shape.
///
/// # Broadcasting Rules
///
/// - Shapes are aligned from the right (trailing dimensions)
/// - Missing dimensions are padded with 1 on the left
/// - For each dimension, sizes must either match or one must be 1
/// - The result dimension is the maximum of the two
///
/// # Examples
///
/// ```ignore
/// // Scalar + Vector: [] + [3] -> [3]
/// let scalar = Tensor::from_slice([5.0f32]);
/// let vector = Tensor::from_slice([1.0f32, 2.0, 3.0]);
/// let (a, b) = scalar.broadcast_for_binop(&vector)?;
///
/// // Matrix + Row: [2, 3] + [1, 3] -> [2, 3]
/// let matrix = Tensor::from_slice([1.0f32; 6]).try_reshape(&[2, 3])?;
/// let row = Tensor::from_slice([1.0f32; 3]).try_reshape(&[1, 3])?;
/// let (a, b) = matrix.broadcast_for_binop(&row)?;
/// ```
///
/// # Errors
///
/// Returns error if shapes are incompatible for broadcasting.
pub(crate) fn broadcast_for_binop(&self, other: &Tensor) -> Result<(Tensor, Tensor)> {
let self_shape = self.shape()?;
let other_shape = other.shape()?;
// Early return if shapes already match
if self_shape == other_shape {
return Ok((self.clone(), other.clone()));
}
// Handle scalar cases (empty shape means scalar in svod)
// Actually, in svod scalars have shape [1], but let's handle both
if self_shape.is_empty() && other_shape.is_empty() {
return Ok((self.clone(), other.clone()));
}
// Align shapes (pad with 1s on left)
let aligned = align_shapes_left(&[self_shape.clone(), other_shape.clone()]);
// Compute broadcast result shape
let result_shape = broadcast_shape(&aligned[0], &aligned[1]).context(UOpSnafu)?;
// Broadcast each tensor to result shape
let self_broadcast = self.broadcast_to(&result_shape)?;
let other_broadcast = other.broadcast_to(&result_shape)?;
Ok((self_broadcast, other_broadcast))
}
/// Broadcast tensor to a target shape.
///
/// This is the low-level broadcast operation that reshapes (adds explicit 1 dimensions)
/// and then expands (replicates data along size-1 dimensions).
///
/// # Algorithm
///
/// 1. If shape already matches, return self
/// 2. Pad shape with 1s on the left to match rank
/// 3. Reshape to add explicit 1 dimensions
/// 4. Expand size-1 dimensions to target size
///
/// # Examples
///
/// ```ignore
/// // [3] -> [2, 3]
/// let t = Tensor::from_slice([1.0f32, 2.0, 3.0]);
/// let target = vec![SInt::from(2), SInt::from(3)];
/// let broadcasted = t.broadcast_to(&target)?;
/// ```
///
/// # Errors
///
/// Returns error if:
/// - Shape has more dimensions than target
/// - Dimension sizes are incompatible (not 1 and not equal to target)
pub fn broadcast_to(&self, target_shape: &svod_ir::shape::Shape) -> Result<Tensor> {
let self_shape = self.shape()?;
// Early return if already correct shape
if &self_shape == target_shape {
return Ok(self.clone());
}
// Cannot broadcast to fewer dimensions
if self_shape.len() > target_shape.len() {
return Err(Error::BroadcastFewerDimensions { from_dims: self_shape.len(), to_dims: target_shape.len() });
}
// Pad shape with 1s on left if needed
let aligned_shape = if self_shape.len() < target_shape.len() {
let padding = target_shape.len() - self_shape.len();
let mut new_shape = svod_ir::shape::Shape::new();
new_shape.extend(std::iter::repeat_n(svod_ir::SInt::from(1), padding));
new_shape.extend(self_shape.iter().cloned());
new_shape
} else {
self_shape.clone()
};
// Validate broadcast compatibility
for (i, (aligned_dim, target_dim)) in aligned_shape.iter().zip(target_shape.iter()).enumerate() {
if let (Some(aligned_size), Some(target_size)) = (aligned_dim.as_const(), target_dim.as_const())
&& aligned_size != 1
&& aligned_size != target_size
{
return Err(Error::BroadcastIncompatible { dim: i, from_size: aligned_size, to_size: target_size });
}
// For symbolic dimensions, conservatively assume they're compatible
}
// Reshape to add explicit 1 dimensions (if needed)
let reshaped = if aligned_shape != self_shape {
// Call IR layer directly to support symbolic dimensions
self.uop().try_reshape(&aligned_shape).map(Self::new).context(UOpSnafu)?
} else {
self.clone()
};
// Check if expansion is actually needed
if &aligned_shape == target_shape {
return Ok(reshaped);
}
// Expand to target shape - call IR layer directly to support symbolic dimensions
reshaped.uop().try_expand(target_shape).map(Self::new).context(UOpSnafu)
}
}