treams_core/numerics/
broadcast.rs

1//! Elementwise evaluation over broadcast inputs, and its pullback. treams-rs extension.
2//!
3//! Each input is a scalar (length one) or has the common output length, as in `NumPy`
4//! broadcasting of one-dimensional arrays. The pullback turns the cotangent of each
5//! output, the gradient of a real loss with respect to that output, into the
6//! cotangents of the inputs. A scalar input's cotangent is the sum of its per-output
7//! cotangents, added in output order on every schedule, so the thread count never
8//! changes it.
9
10use rayon::prelude::*;
11
12use super::parallel::Parallel;
13use crate::{Complex, Error, Result};
14
15/// Common output length of scalar-or-array inputs: zero when an input is empty,
16/// otherwise the longest length. Any other length is rejected with `message`.
17pub(crate) fn size(lengths: &[usize], message: &str) -> Result<usize> {
18    let size = if lengths.contains(&0) {
19        0
20    } else {
21        lengths.iter().copied().max().unwrap_or_default()
22    };
23    if lengths.iter().any(|&n| n != 1 && n != size) {
24        return Err(Error::InvalidInput(message.into()));
25    }
26    Ok(size)
27}
28
29/// Element `i` of a scalar-or-array input of validated length.
30#[allow(clippy::indexing_slicing)] // Scalar-or-element indexing after `size` validation.
31pub(crate) fn element<T: Copy>(values: &[T], i: usize) -> T {
32    values[if values.len() == 1 { 0 } else { i }]
33}
34
35/// Evaluate `f` at `0..size` in order, stopping at an error.
36pub(crate) fn map<T: Send>(
37    size: usize,
38    parallel: Parallel,
39    f: impl Fn(usize) -> Result<T> + Sync + Send,
40) -> Result<Vec<T>> {
41    match parallel.chunk(size) {
42        Some(chunk) => crate::threads::install(|| {
43            (0..size)
44                .into_par_iter()
45                .with_min_len(chunk)
46                .map(f)
47                .collect()
48        }),
49        None => (0..size).map(f).collect(),
50    }
51}
52
53/// Apply one input direction to an elementwise evaluation. Each tangent is scalar
54/// or has one finite entry per output; `f` receives the direction at that output.
55pub(crate) fn pushforward<const A: usize, T: Send>(
56    tangents: [&[Complex]; A],
57    size: usize,
58    message: &str,
59    parallel: Parallel,
60    f: impl Fn(usize, [Complex; A]) -> Result<T> + Sync + Send,
61) -> Result<Vec<T>> {
62    if tangents.iter().any(|tangent| {
63        (tangent.len() != 1 && tangent.len() != size)
64            || tangent.iter().any(|&value| !crate::numerics::finite(value))
65    }) {
66        return Err(Error::InvalidInput(message.into()));
67    }
68    map(size, parallel, |i| {
69        f(i, tangents.map(|tangent| element(tangent, i)))
70    })
71}
72
73/// The cotangent of one output element.
74pub(crate) trait Cotangent: Sync {
75    /// Whether every component is finite.
76    fn finite(&self) -> bool;
77}
78
79impl Cotangent for Complex {
80    fn finite(&self) -> bool {
81        crate::numerics::finite(*self)
82    }
83}
84
85impl<const N: usize> Cotangent for [Complex; N] {
86    fn finite(&self) -> bool {
87        self.iter().all(|&z| crate::numerics::finite(z))
88    }
89}
90
91/// Pull back an elementwise evaluation with `A` broadcast complex arguments of the
92/// given `lengths`. The cotangent needs one finite entry per output (otherwise
93/// `message` is returned); `f(i, g)` gives the argument cotangents of output `i`.
94pub(crate) fn pullback<G: Cotangent, const A: usize>(
95    cotangent: &[G],
96    size: usize,
97    message: &str,
98    lengths: [usize; A],
99    parallel: Parallel,
100    f: impl Fn(usize, &G) -> Result<[Complex; A]> + Sync + Send,
101) -> Result<[Vec<Complex>; A]> {
102    if cotangent.len() != size || !cotangent.iter().all(Cotangent::finite) {
103        return Err(Error::InvalidInput(message.into()));
104    }
105    let Some(chunk) = parallel.chunk(size) else {
106        // In order on this thread: gather each output's argument cotangents
107        // straight into the gradients, without per-output intermediates. A
108        // scalar argument's sum starts from zero and adds in output order, as
109        // `Iterator::sum` does below.
110        let mut sums = [Complex::default(); A];
111        let mut gradients =
112            lengths.map(|length| Vec::with_capacity(if length == 1 { 1 } else { size }));
113        for (i, g) in cotangent.iter().enumerate() {
114            let values = f(i, g)?;
115            for (((gradient, sum), &length), value) in gradients
116                .iter_mut()
117                .zip(&mut sums)
118                .zip(&lengths)
119                .zip(values)
120            {
121                if length == 1 {
122                    *sum += value;
123                } else {
124                    gradient.push(value);
125                }
126            }
127        }
128        for ((gradient, sum), &length) in gradients.iter_mut().zip(sums).zip(&lengths) {
129            if length == 1 {
130                gradient.push(sum);
131            }
132        }
133        return Ok(gradients);
134    };
135    // In parallel: write each output's argument cotangents into its row in
136    // place, which indexed iterators split without collecting partial vectors.
137    let mut rows = vec![[Complex::default(); A]; size];
138    crate::threads::install(|| {
139        rows.par_iter_mut()
140            .zip(cotangent)
141            .enumerate()
142            .with_min_len(chunk)
143            .try_for_each(|(i, (row, g))| {
144                *row = f(i, g)?;
145                Ok(())
146            })
147    })?;
148    if A == 1 && lengths != [1; A] {
149        // The rows of a single array argument are its gradient already.
150        let mut gradient = Some(rows.into_flattened());
151        return Ok(std::array::from_fn(|_| gradient.take().unwrap_or_default()));
152    }
153    // A scalar argument's terms add in output order, as on the calling thread.
154    Ok(std::array::from_fn(|a| {
155        let gradient = rows
156            .iter()
157            .map(|row| row.get(a).copied().unwrap_or_default());
158        if lengths.get(a) == Some(&1) {
159            vec![gradient.sum()]
160        } else {
161            gradient.collect()
162        }
163    }))
164}
165
166/// Broadcast lengths, output order and the pullback on every schedule.
167#[cfg(test)]
168mod tests {
169    use proptest::{prelude::*, test_runner::TestCaseError};
170
171    use super::{Parallel, element, map, pullback, pushforward, size};
172    use crate::{Complex, test_support::DEFAULT_CASES};
173
174    #[test]
175    fn size_accepts_scalars_and_one_common_length() {
176        let message = "arrays must have equal lengths or scalar inputs";
177        for (lengths, expected) in [
178            (&[1, 1][..], Some(1)),
179            (&[1, 4, 4], Some(4)),
180            (&[4, 1, 1], Some(4)),
181            (&[0, 1], Some(0)),
182            (&[0], Some(0)),
183            (&[], Some(0)),
184            (&[0, 3], None),
185            (&[2, 3], None),
186            (&[2, 1, 3], None),
187        ] {
188            match size(lengths, message) {
189                Ok(size) => assert_eq!(Some(size), expected, "{lengths:?}"),
190                Err(error) => {
191                    assert_eq!(expected, None, "{lengths:?}");
192                    assert_eq!(error.to_string(), message);
193                }
194            }
195        }
196    }
197
198    proptest! {
199        #![proptest_config(ProptestConfig::with_cases(DEFAULT_CASES))]
200
201        #[test]
202        fn pullback_is_the_transposed_elementwise_product(
203            size in 0_usize..40,
204            scalar in any::<[bool; 2]>(),
205            threshold in 0_usize..48,
206            chunked in any::<bool>(),
207        ) {
208            check_pullback(size, scalar, threshold, chunked)?;
209        }
210    }
211
212    /// For `f_i(a, b) = a_i b_i` over broadcast arguments, the evaluation keeps output
213    /// order under every schedule, and the pullback pairs each cotangent with the
214    /// other argument and sums it for scalar arguments, bit for bit alike in order
215    /// on the calling thread and in parallel.
216    fn check_pullback(
217        size: usize,
218        scalar: [bool; 2],
219        threshold: usize,
220        chunked: bool,
221    ) -> Result<(), TestCaseError> {
222        let parallel = if chunked {
223            Parallel::Chunked(threshold)
224        } else {
225            Parallel::AtLeast(threshold)
226        };
227        let argument = |index: usize, scalar: bool| -> Vec<Complex> {
228            let length = if scalar { 1 } else { size };
229            (0..length)
230                .map(|i| Complex::new(f64::from(u32::try_from(i + index).unwrap()), 0.5))
231                .collect()
232        };
233        let arguments = [argument(1, scalar[0]), argument(3, scalar[1])];
234        let lengths = arguments.each_ref().map(Vec::len);
235        let values = map(size, parallel, |i| {
236            Ok(element(&arguments[0], i) * element(&arguments[1], i))
237        })?;
238        prop_assert_eq!(values.len(), size);
239        for (i, value) in values.iter().enumerate() {
240            prop_assert_eq!(
241                *value,
242                element(&arguments[0], i) * element(&arguments[1], i)
243            );
244        }
245        let cotangent: Vec<_> = (0..size)
246            .map(|i| Complex::new(0.25, f64::from(u32::try_from(i).unwrap())))
247            .collect();
248        let tangents = arguments.each_ref().map(|values| {
249            values
250                .iter()
251                .map(|z| Complex::new(0.3, -0.2) * z)
252                .collect::<Vec<_>>()
253        });
254        let direction = pushforward(
255            tangents.each_ref().map(Vec::as_slice),
256            size,
257            "tangent",
258            parallel,
259            |i, [da, db]| Ok(da * element(&arguments[1], i) + element(&arguments[0], i) * db),
260        )?;
261        for (i, &actual) in direction.iter().enumerate() {
262            let expected = element(&tangents[0], i) * element(&arguments[1], i)
263                + element(&arguments[0], i) * element(&tangents[1], i);
264            prop_assert_eq!(actual, expected);
265        }
266        let gradients = pullback(&cotangent, size, "cotangent", lengths, parallel, |i, &g| {
267            let [a, b] = arguments.each_ref().map(|v| element(v, i));
268            Ok([g * b.conj(), g * a.conj()])
269        })?;
270        for (argument, gradient) in arguments.iter().zip(&gradients) {
271            prop_assert_eq!(gradient.len(), argument.len());
272        }
273        for (index, gradient) in gradients.iter().enumerate() {
274            let other = &arguments[1 - index];
275            let expected: Vec<_> = (0..size)
276                .map(|i| cotangent[i] * element(other, i).conj())
277                .collect();
278            let expected = if lengths[index] == 1 {
279                vec![expected.iter().sum()]
280            } else {
281                expected
282            };
283            // Every schedule adds a scalar argument's terms in output order.
284            prop_assert_eq!(gradient, &expected);
285        }
286        // A single argument takes its own path when it is an array.
287        let [single] = pullback(
288            &cotangent,
289            size,
290            "cotangent",
291            [lengths[0]],
292            parallel,
293            |i, &g| Ok([g * element(&arguments[1], i).conj()]),
294        )?;
295        prop_assert_eq!(&single, &gradients[0]);
296        Ok(())
297    }
298
299    #[test]
300    fn pullback_rejects_mismatched_or_nonfinite_cotangents() {
301        let zero = |_: usize, _: &[Complex; 2]| Ok([Complex::default(); 1]);
302        let finite = [[Complex::new(0.5, 0.0); 2]; 2];
303        let parallel = Parallel::AtLeast(1024);
304        assert!(pullback(&finite, 2, "cotangent", [2], parallel, zero).is_ok());
305        assert!(pullback(&finite, 3, "cotangent", [3], parallel, zero).is_err());
306        assert!(pullback(&finite[..1], 2, "cotangent", [2], parallel, zero).is_err());
307        let mut nonfinite = finite;
308        nonfinite[1][0].im = f64::NAN;
309        let error = pullback(&nonfinite, 2, "cotangent", [2], parallel, zero).unwrap_err();
310        assert_eq!(error.to_string(), "cotangent");
311    }
312
313    #[test]
314    fn pushforward_rejects_mismatched_or_nonfinite_tangents() {
315        let parallel = Parallel::AtLeast(1024);
316        let identity = |_: usize, [value]: [Complex; 1]| Ok(value);
317        let finite = [Complex::new(0.5, 0.2); 2];
318        assert!(pushforward([&finite], 2, "tangent", parallel, identity).is_ok());
319        assert!(pushforward([&finite[..1]], 2, "tangent", parallel, identity).is_ok());
320        assert!(pushforward([&finite], 3, "tangent", parallel, identity).is_err());
321        assert!(pushforward([&[]], 2, "tangent", parallel, identity).is_err());
322        assert!(pushforward([&[]], 0, "tangent", parallel, identity).is_ok());
323        let bad = [Complex::new(f64::NAN, 0.0)];
324        assert!(pushforward([&bad], 2, "tangent", parallel, identity).is_err());
325    }
326}
treams_core/numerics/broadcast.rs

treams_core/numerics/
broadcast.rs
