Skip to main content

diffsol_la/matrix/
mod.rs

1use std::fmt::Debug;
2use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
3
4use crate::error::LaError;
5use crate::scalar::Scale;
6use crate::{Context, IndexType, Scalar, Vector, VectorIndex};
7
8use extract_block::combine;
9use sparsity::{Dense, MatrixSparsity, MatrixSparsityRef};
10
11#[cfg(feature = "cuda")]
12pub mod cuda;
13
14#[cfg(feature = "cuda-oxide")]
15pub mod cuda_oxide;
16
17#[cfg(feature = "nalgebra")]
18pub mod dense_nalgebra_serial;
19
20#[cfg(feature = "faer")]
21pub mod dense_faer_serial;
22
23#[cfg(feature = "faer")]
24pub mod sparse_faer;
25
26pub mod default_solver;
27pub mod extract_block;
28pub mod sparsity;
29
30/// Common interface for matrix types, providing access to scalar type, context, and dimensions.
31pub trait MatrixCommon: Sized + Debug {
32    type V: Vector<T = Self::T, C = Self::C, Index: VectorIndex<C = Self::C>>;
33    type T: Scalar;
34    type C: Context;
35    type Inner;
36
37    /// Get the number of rows in this matrix.
38    fn nrows(&self) -> IndexType;
39    /// Get the number of columns in this matrix.
40    fn ncols(&self) -> IndexType;
41    /// Get a reference to the inner representation of the matrix.
42    fn inner(&self) -> &Self::Inner;
43}
44
45impl<M> MatrixCommon for &M
46where
47    M: MatrixCommon,
48{
49    type T = M::T;
50    type V = M::V;
51    type C = M::C;
52    type Inner = M::Inner;
53
54    fn nrows(&self) -> IndexType {
55        M::nrows(*self)
56    }
57    fn ncols(&self) -> IndexType {
58        M::ncols(*self)
59    }
60    fn inner(&self) -> &Self::Inner {
61        M::inner(*self)
62    }
63}
64
65impl<M> MatrixCommon for &mut M
66where
67    M: MatrixCommon,
68{
69    type T = M::T;
70    type V = M::V;
71    type C = M::C;
72    type Inner = M::Inner;
73
74    fn ncols(&self) -> IndexType {
75        M::ncols(*self)
76    }
77    fn nrows(&self) -> IndexType {
78        M::nrows(*self)
79    }
80    fn inner(&self) -> &Self::Inner {
81        M::inner(*self)
82    }
83}
84
85/// Operations on matrices by value (addition and subtraction).
86///
87/// This trait defines matrix addition and subtraction when both operands are owned or references.
88pub trait MatrixOpsByValue<Rhs = Self, Output = Self>:
89    MatrixCommon + Add<Rhs, Output = Output> + Sub<Rhs, Output = Output>
90{
91}
92
93impl<M, Rhs, Output> MatrixOpsByValue<Rhs, Output> for M where
94    M: MatrixCommon + Add<Rhs, Output = Output> + Sub<Rhs, Output = Output>
95{
96}
97
98/// In-place operations on matrices (addition and subtraction).
99///
100/// This trait defines in-place matrix addition and subtraction (self += rhs, self -= rhs).
101pub trait MatrixMutOpsByValue<Rhs = Self>: MatrixCommon + AddAssign<Rhs> + SubAssign<Rhs> {}
102
103impl<M, Rhs> MatrixMutOpsByValue<Rhs> for M where M: MatrixCommon + AddAssign<Rhs> + SubAssign<Rhs> {}
104
105/// A trait allowing for references to implement matrix operations
106pub trait MatrixRef<M: MatrixCommon>: Mul<Scale<M::T>, Output = M> {}
107impl<RefT, M: MatrixCommon> MatrixRef<M> for RefT where RefT: Mul<Scale<M::T>, Output = M> {}
108
109/// A base matrix trait supporting both sparse and dense matrices.
110///
111/// This trait provides a complete interface for matrix operations including:
112/// - Matrix creation and memory management
113/// - Matrix-vector and matrix-matrix multiplication
114/// - Element access and modification
115/// - Sparsity information and handling
116/// - Matrix decomposition and combination operations
117/// - Triplet-based construction for sparse matrices
118///
119/// Implementing matrices can be dense or sparse, and may be hosted on CPU or GPU.
120/// Users typically do not need to implement this trait; use provided implementations.
121pub trait Matrix:
122    MatrixCommon + Mul<Scale<Self::T>, Output = Self> + Clone + Send + 'static
123{
124    type Sparsity: MatrixSparsity<Self>;
125    type SparsityRef<'a>: MatrixSparsityRef<'a, Self>
126    where
127        Self: 'a;
128
129    /// Return sparsity information, or `None` if the matrix is dense.
130    fn sparsity(&self) -> Option<Self::SparsityRef<'_>>;
131
132    /// Get the context associated with this matrix (for device placement, memory management, etc.).
133    fn context(&self) -> &Self::C;
134
135    /// Get a mutable reference to the inner representation of the matrix.
136    fn inner_mut(&mut self) -> &mut Self::Inner;
137
138    /// Returns true if this matrix is stored in a sparse format
139    fn is_sparse() -> bool {
140        Self::zeros(1, 1, Default::default()).sparsity().is_some()
141    }
142
143    /// Partition the diagonal indices into two groups: those with zero diagonal elements and those with non-zero diagonal elements.
144    ///
145    /// This is useful for identifying algebraic constraints, which typically have zero diagonal elements in the mass matrix.
146    /// Returns a tuple of (zero_diagonal_indices, non_zero_diagonal_indices).
147    /// For batched matrices the partition is taken from batch 0, so all batches are assumed
148    /// to share the same structure.
149    fn partition_indices_by_zero_diagonal(
150        &self,
151    ) -> (<Self::V as Vector>::Index, <Self::V as Vector>::Index);
152
153    /// Perform a matrix-vector multiplication: y = alpha * self * x + beta * y
154    fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V);
155
156    /// Copy the contents of `other` into this matrix.
157    fn copy_from(&mut self, other: &Self);
158
159    /// Create a new matrix of shape `nrows` x `ncols` filled with zeros.
160    fn zeros(nrows: IndexType, ncols: IndexType, ctx: Self::C) -> Self;
161
162    /// Create a new matrix from a sparsity pattern. Non-zero elements are not initialized.
163    fn new_from_sparsity(
164        nrows: IndexType,
165        ncols: IndexType,
166        sparsity: Option<Self::Sparsity>,
167        ctx: Self::C,
168    ) -> Self;
169
170    /// Create a new diagonal matrix from a vector holding the diagonal elements.
171    fn from_diagonal(v: &Self::V) -> Self;
172
173    /// Set the values of column `j` to be equal to the values in `v`.
174    ///
175    /// For sparse matrices, only the existing non-zero elements are updated.
176    fn set_column(&mut self, j: IndexType, v: &Self::V);
177
178    /// Add a column of this matrix to a vector: v += self[:, j]
179    fn add_column_to_vector(&self, j: IndexType, v: &mut Self::V);
180
181    /// Add every column of this matrix to the matching batch lane of `v`: lane `b * ncols() + j`
182    /// receives column `j` of this matrix's batch `b`.
183    fn add_columns_to_batched_vector(&self, v: &mut Self::V);
184
185    /// Assign the values in the `data` vector to this matrix at the indices in `dst_indices`
186    /// from the indices in `src_indices`.
187    ///
188    /// For dense matrices, the index is the data index in column-major order.
189    /// For sparse matrices, the index is the index into the data array.
190    fn set_data_with_indices(
191        &mut self,
192        dst_indices: &<Self::V as Vector>::Index,
193        src_indices: &<Self::V as Vector>::Index,
194        data: &Self::V,
195    );
196
197    /// Gather values from another matrix at specified indices into this matrix.
198    ///
199    /// For sparse matrices: the index `idx_i` in `indices` is an index into the data array for `other`,
200    /// and is copied to the index `idx_i` in the data array for this matrix.
201    /// For dense matrices: the index is the data index in column-major order.
202    fn gather(&mut self, other: &Self, indices: &<Self::V as Vector>::Index);
203
204    /// Split this matrix into four submatrices based on algebraic constraint indices.
205    ///
206    /// Partitions the matrix into blocks:
207    /// ```text
208    /// M = [UL, UR]
209    ///     [LL, LR]
210    /// ```
211    /// where:
212    /// - UL contains rows and columns NOT in `algebraic_indices`
213    /// - UR contains rows NOT in `algebraic_indices` and columns in `algebraic_indices`
214    /// - LL contains rows in `algebraic_indices` and columns NOT in `algebraic_indices`
215    /// - LR contains rows and columns in `algebraic_indices`
216    ///
217    /// Returns an array of tuples, where each tuple contains a submatrix and the indices that were used to create it.
218    /// These indices can be used with `gather()` to update the submatrix.
219    fn split(
220        &self,
221        algebraic_indices: &<Self::V as Vector>::Index,
222    ) -> [(Self, <Self::V as Vector>::Index); 4] {
223        match self.sparsity() {
224            Some(sp) => sp.split(algebraic_indices).map(|(sp, src_indices)| {
225                let mut m = Self::new_from_sparsity(
226                    sp.nrows(),
227                    sp.ncols(),
228                    Some(sp),
229                    self.context().clone(),
230                );
231                m.gather(self, &src_indices);
232                (m, src_indices)
233            }),
234            None => Dense::<Self>::new(self.nrows(), self.ncols())
235                .split(algebraic_indices)
236                .map(|(sp, src_indices)| {
237                    let mut m = Self::new_from_sparsity(
238                        sp.nrows(),
239                        sp.ncols(),
240                        None,
241                        self.context().clone(),
242                    );
243                    m.gather(self, &src_indices);
244                    (m, src_indices)
245                }),
246        }
247    }
248
249    /// Combine four submatrices back into a single matrix based on algebraic constraint indices.
250    ///
251    /// Inverse operation of `split()`. Takes submatrices `ul`, `ur`, `ll`, `lr` and combines them
252    /// back into the original matrix structure.
253    fn combine(
254        ul: &Self,
255        ur: &Self,
256        ll: &Self,
257        lr: &Self,
258        algebraic_indices: &<Self::V as Vector>::Index,
259    ) -> Self {
260        combine(ul, ur, ll, lr, algebraic_indices)
261    }
262
263    /// Perform the assignment: self = x + beta * y where x and y are matrices and beta is a scalar.
264    ///
265    /// Note: Panics if the sparsity patterns of self, x, and y do not match.
266    /// The sparsity of self must be the union of the sparsity of x and y.
267    fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self);
268
269    /// Iterate over structural positions and values of the matrix.
270    ///
271    /// Returns a tuple:
272    /// - First iterator: `(row, col)` pairs for each non-zero element (length `nnz`)
273    /// - Second iterator: values (length `nnz * nbatch`), laid out batch-contiguously:
274    ///   `[batch0_val0..batch0_valN, batch1_val0..batch1_valN, ...]`
275    fn triplet_iter(
276        &self,
277    ) -> (
278        impl Iterator<Item = (IndexType, IndexType)> + '_,
279        impl Iterator<Item = Self::T> + '_,
280    );
281
282    /// Create a new matrix from structural indices and values.
283    ///
284    /// - `indices`: `(row, col)` pairs for each non-zero element (length `nnz`)
285    /// - `values`: values laid out batch-contiguously (length `nnz * ctx.nbatch()`):
286    ///   `[batch0_val0..batch0_valN, batch1_val0..batch1_valN, ...]`
287    fn try_from_triplets(
288        nrows: IndexType,
289        ncols: IndexType,
290        indices: Vec<(IndexType, IndexType)>,
291        values: Vec<Self::T>,
292        ctx: Self::C,
293    ) -> Result<Self, LaError>;
294}
295
296/// Largest column range the small-coefficient kernels accept — [`DenseMatrix::mul_cols_by`] and
297/// [`DenseMatrix::gemv_cols`].
298///
299/// Both take their coefficients from the host: the serial backends stage them in a fixed stack
300/// buffer and the CUDA backend passes them to the kernel by value, so neither allocates — which
301/// needs a bound. Two things drive it: the BDF difference table is at most `MAX_ORDER + 1 = 6`
302/// columns wide, and the widest built-in Runge-Kutta tableau (`tsit45`) has 7 stages. 8 covers
303/// both.
304///
305/// Must match `MAX_SMALL_COLS` in `src/cuda_kernels/cuda_kernels_common.h`, which sizes the
306/// kernels' by-value coefficient arrays; nothing checks that across the language boundary.
307pub const MAX_SMALL_COLS: usize = 8;
308
309/// A dense column-major matrix with efficient column access operations.
310///
311/// This trait represents matrices stored in column-major order, where accessing matrix columns
312/// is efficient. It supports:
313/// - Column operations (access, modification, and the fused column kernels below)
314/// - Element access and modification
315/// - Matrix resizing
316///
317/// The column-major layout makes operations on individual or ranges of columns very efficient.
318pub trait DenseMatrix:
319    Matrix + for<'b> MatrixOpsByValue<&'b Self, Self> + for<'b> MatrixMutOpsByValue<&'b Self>
320{
321    /// Update a table of backward differences held in columns `0..=order+2` of
322    /// this matrix, given the newly computed highest-order correction `d`.
323    ///
324    /// Equivalent to:
325    /// ```text
326    /// self[:, order+2] = d - self[:, order+1]
327    /// for i in (0..=order+1).rev() { self[:, i] += self[:, i+1] }
328    /// ```
329    /// where each addition uses the just-updated value of column `i+1` (a
330    /// right-to-left cumulative scan).
331    fn update_backward_diff(&mut self, order: IndexType, d: &Self::V);
332
333    /// Right-multiply a range of columns by a small matrix, in place:
334    ///
335    /// ```text
336    /// self[:, 0..ncols] = self[:, 0..ncols] * rhs
337    /// ```
338    ///
339    /// `rhs` is an `ncols x ncols` block held column-major in a host slice of exactly
340    /// `ncols * ncols` values — it is small (the BDF step-size update uses at most `6 x 6`).
341    ///
342    /// Panics if `ncols` exceeds [`MAX_SMALL_COLS`], since the serial backends stage the old
343    /// column values in a fixed-size buffer rather than allocating.
344    fn mul_cols_by(&mut self, ncols: IndexType, rhs: &[Self::T]);
345
346    /// Matrix-vector multiply over a range of columns, with host-side coefficients:
347    ///
348    /// ```text
349    /// y = alpha * self[:, start..end] * x + beta * y
350    /// ```
351    ///
352    /// `x` is a host slice, not a vector, because `x.len()` is the *width of the column range*
353    /// rather than a dimension of the problem: every caller in the solvers passes a short
354    /// coefficient list computed from host scalars — a Runge-Kutta tableau row, the `b`/`d`
355    /// vectors, interpolation weights, BDF's psi weights. Keeping them on the host is what lets
356    /// those callers build them in a stack array instead of allocating a vector per call.
357    ///
358    /// Use [`Matrix::gemv`] instead when `x` is a genuine operand of length `ncols` — a
359    /// Jacobian- or mass-times-vector product.
360    ///
361    /// Only the leading `end - start` columns take part; the rest of the matrix is not read.
362    /// `x.len()` must be at least `end - start`; only its leading `end - start` entries are
363    /// read, as with a BLAS `gemv` given a longer vector and `incx = 1`. `beta == 0` means `y`
364    /// is not read, so it may hold uninitialised values. An empty range contributes nothing,
365    /// leaving `y = beta * y`.
366    ///
367    /// The coefficients are shared by every batch — unlike [`Matrix::gemv`], this cannot apply
368    /// a different `x` per batch. Panics if `end - start` exceeds [`MAX_SMALL_COLS`].
369    fn gemv_cols(
370        &self,
371        start: IndexType,
372        end: IndexType,
373        alpha: Self::T,
374        x: &[Self::T],
375        beta: Self::T,
376        y: &mut Self::V,
377    );
378
379    /// Get an immutable vector view of column `i`.
380    fn column(&self, i: IndexType) -> <Self::V as Vector>::View<'_>;
381
382    /// Get a mutable vector view of column `i`.
383    fn column_mut(&mut self, i: IndexType) -> <Self::V as Vector>::ViewMut<'_>;
384
385    /// Set the value at the given row and column indices.
386    ///
387    /// Scalar access is single-system, as it is for [`Vector::set_index`]: panics unless
388    /// `nbatch == 1`. So does `IndexMut`.
389    fn set_index(&mut self, i: IndexType, j: IndexType, value: Self::T);
390
391    /// Set the value at the given row and column indices, in a single batch.
392    ///
393    /// The batch-aware counterpart to [`Self::set_index`]: `(i, j)` alone is ambiguous once
394    /// `nbatch > 1`, so the batch is named explicitly. Mirrors [`Vector::get_batch_mut`].
395    fn set_index_batch(&mut self, batch: IndexType, i: IndexType, j: IndexType, value: Self::T);
396
397    /// Get the value at the given row and column indices.
398    ///
399    /// Scalar access is single-system, as it is for [`Vector::get_index`]: panics unless
400    /// `nbatch == 1`. So does `Index`.
401    fn get_index(&self, i: IndexType, j: IndexType) -> Self::T;
402
403    /// Resize the number of columns in the matrix, preserving existing data.
404    ///
405    /// New elements (if added) are uninitialized. If the number of columns decreases, trailing columns are discarded.
406    fn resize_cols(&mut self, ncols: IndexType);
407
408    /// Create a new matrix from a vector of values in column-major order.
409    ///
410    /// The values are assumed to be stored in column-major order (first column, then second column, etc.).
411    fn from_vec(nrows: IndexType, ncols: IndexType, data: Vec<Self::T>, ctx: Self::C) -> Self;
412}
413
414#[cfg(test)]
415#[allow(clippy::items_after_test_module)]
416pub(crate) mod tests {
417    use std::ops::{Add, Mul, Sub};
418
419    use super::{
420        DenseMatrix, Matrix, MatrixCommon, MatrixSparsity, MatrixSparsityRef, MAX_SMALL_COLS,
421    };
422    use crate::scalar::Scale;
423    use crate::{scalar::IndexType, Context, Vector, VectorIndex, VectorView, VectorViewMut};
424    use num_traits::{FromPrimitive, One, Zero};
425
426    fn f<M: Matrix>(x: f64) -> M::T {
427        M::T::from_f64(x).unwrap()
428    }
429
430    /// Read one element of one batch. `get_index` is single-system, so a batched matrix is read a
431    /// column at a time.
432    fn at<M: DenseMatrix>(a: &M, batch: usize, i: IndexType, j: IndexType) -> M::T {
433        a.column(j).into_owned().get_batch(batch).get_index(i)
434    }
435
436    pub(crate) fn triplet_values<M: Matrix>(m: &M) -> Vec<M::T> {
437        let (_, vals) = m.triplet_iter();
438        vals.collect()
439    }
440
441    fn triplet_indices<M: Matrix>(m: &M) -> Vec<(IndexType, IndexType)> {
442        let (idx, _) = m.triplet_iter();
443        idx.collect()
444    }
445
446    pub fn test_partition_indices_by_zero_diagonal<M: Matrix>() {
447        let indices = vec![(0, 0), (1, 1), (3, 3)];
448        let values = vec![M::T::one(), M::T::from_f64(2.0).unwrap(), M::T::one()];
449        let m = M::try_from_triplets(4, 4, indices, values, Default::default()).unwrap();
450        let (zero_diagonal_indices, non_zero_diagonal_indices) =
451            m.partition_indices_by_zero_diagonal();
452        assert_eq!(zero_diagonal_indices.clone_as_vec(), vec![2]);
453        assert_eq!(non_zero_diagonal_indices.clone_as_vec(), vec![0, 1, 3]);
454
455        let indices = vec![(0, 0), (1, 1), (2, 2), (3, 3)];
456        let values = vec![
457            M::T::one(),
458            M::T::from_f64(2.0).unwrap(),
459            M::T::zero(),
460            M::T::one(),
461        ];
462        let m = M::try_from_triplets(4, 4, indices, values, Default::default()).unwrap();
463        let (zero_diagonal_indices, non_zero_diagonal_indices) =
464            m.partition_indices_by_zero_diagonal();
465        assert_eq!(zero_diagonal_indices.clone_as_vec(), vec![2]);
466        assert_eq!(non_zero_diagonal_indices.clone_as_vec(), vec![0, 1, 3]);
467
468        let indices = vec![(0, 0), (1, 1), (2, 2), (3, 3)];
469        let values = vec![
470            M::T::one(),
471            M::T::from_f64(2.0).unwrap(),
472            M::T::from_f64(3.0).unwrap(),
473            M::T::one(),
474        ];
475        let m = M::try_from_triplets(4, 4, indices, values, Default::default()).unwrap();
476        let (zero_diagonal_indices, non_zero_diagonal_indices) =
477            m.partition_indices_by_zero_diagonal();
478        assert_eq!(
479            zero_diagonal_indices.clone_as_vec(),
480            Vec::<IndexType>::new()
481        );
482        assert_eq!(non_zero_diagonal_indices.clone_as_vec(), vec![0, 1, 2, 3]);
483
484        // column 2 has a single stored entry below the diagonal (row 3) and none at (2, 2):
485        // the sparse backend's early-exit ("row index has passed the diagonal") branch fires.
486        let indices = vec![(3, 2)];
487        let values = vec![M::T::one()];
488        let m = M::try_from_triplets(4, 4, indices, values, Default::default()).unwrap();
489        let (zero_diagonal_indices, non_zero_diagonal_indices) =
490            m.partition_indices_by_zero_diagonal();
491        assert_eq!(zero_diagonal_indices.clone_as_vec(), vec![0, 1, 2, 3]);
492        assert_eq!(
493            non_zero_diagonal_indices.clone_as_vec(),
494            Vec::<IndexType>::new()
495        );
496    }
497
498    // --- Matrix-generic tests (work with both dense and sparse) ---
499
500    pub fn test_zeros<M: Matrix>() {
501        let a = M::zeros(2, 3, Default::default());
502        assert_eq!(a.nrows(), 2);
503        assert_eq!(a.ncols(), 3);
504        let vals = triplet_values(&a);
505        assert!(vals.is_empty() || vals.iter().all(|v| v.is_zero()));
506        assert_eq!(M::is_sparse(), a.sparsity().is_some());
507    }
508
509    /// `M::Sparsity`/`M::SparsityRef<'_>` trait methods (`nrows`/`ncols`/`is_sparse`/`indices`/
510    /// `union`/`new_diagonal`/`get_index`/`as_ref`/`to_owned`/`split`): only `is_sparse` was ever
511    /// called generically before, via `test_zeros`, so none of the sparsity pattern machinery
512    /// itself (used by both the `Dense` stand-in for dense backends and the real sparse pattern)
513    /// was exercised.
514    pub fn test_sparsity<M: Matrix>() {
515        let a = M::Sparsity::try_from_indices(2, 2, vec![(0, 0), (1, 1)]).unwrap();
516        assert_eq!(MatrixSparsity::<M>::nrows(&a), 2);
517        assert_eq!(MatrixSparsity::<M>::ncols(&a), 2);
518        assert_eq!(
519            M::is_sparse(),
520            <M::Sparsity as MatrixSparsity<M>>::is_sparse()
521        );
522        let _ = a.indices();
523        let diag = M::Sparsity::new_diagonal(3);
524        assert_eq!(MatrixSparsity::<M>::nrows(&diag), 3);
525        assert_eq!(MatrixSparsity::<M>::ncols(&diag), 3);
526
527        let b = M::Sparsity::try_from_indices(2, 2, vec![(0, 1)]).unwrap();
528        let unioned = a.clone().union(b.as_ref()).unwrap();
529        assert_eq!(MatrixSparsity::<M>::nrows(&unioned), 2);
530        assert_eq!(MatrixSparsity::<M>::ncols(&unioned), 2);
531
532        let idx = a.get_index(&[(0, 0), (1, 1)], Default::default());
533        assert_eq!(idx.len(), 2);
534
535        let a_ref = a.as_ref();
536        assert_eq!(MatrixSparsityRef::<M>::nrows(&a_ref), 2);
537        assert_eq!(MatrixSparsityRef::<M>::ncols(&a_ref), 2);
538        assert_eq!(
539            M::is_sparse(),
540            <M::SparsityRef<'_> as MatrixSparsityRef<M>>::is_sparse()
541        );
542        let _ = a_ref.indices();
543        let owned = a_ref.to_owned();
544        assert_eq!(MatrixSparsity::<M>::nrows(&owned), 2);
545
546        let algebraic = <M::V as Vector>::Index::from_vec(vec![1], Default::default());
547        let [(ul, _), (_, _), (_, _), (lr, _)] = a_ref.split(&algebraic);
548        assert_eq!(MatrixSparsity::<M>::nrows(&ul), 1);
549        assert_eq!(MatrixSparsity::<M>::nrows(&lr), 1);
550    }
551
552    pub fn test_matrix_common_by_ref<M: Matrix>() {
553        fn via_common<T: MatrixCommon>(m: T) -> (IndexType, IndexType) {
554            let dims = (m.nrows(), m.ncols());
555            let _ = m.inner();
556            dims
557        }
558        let mut a = M::zeros(2, 3, Default::default());
559        assert_eq!(via_common(&a), (2, 3));
560        assert_eq!(via_common(&mut a), (2, 3));
561    }
562
563    pub fn test_from_diagonal<M: Matrix>() {
564        let v = M::V::from_vec(
565            vec![f::<M>(2.0), f::<M>(3.0), f::<M>(5.0)],
566            Default::default(),
567        );
568        let a = M::from_diagonal(&v);
569        assert_eq!(a.nrows(), 3);
570        assert_eq!(a.ncols(), 3);
571        let idx = triplet_indices(&a);
572        let vals = triplet_values(&a);
573        // diagonal matrix triplet_iter returns only the diagonal nnz entries
574        for &(i, j) in &idx {
575            let pos = idx.iter().position(|&x| x == (i, j)).unwrap();
576            if i == j {
577                assert!(
578                    vals[pos] != M::T::zero(),
579                    "diagonal entry should be non-zero"
580                );
581            } else {
582                assert!(vals[pos].is_zero(), "off-diagonal entry should be zero");
583            }
584        }
585    }
586
587    pub fn test_from_diagonal_dense<M: DenseMatrix>() {
588        let v = M::V::from_vec(
589            vec![f::<M>(2.0), f::<M>(3.0), f::<M>(5.0)],
590            Default::default(),
591        );
592        let a = M::from_diagonal(&v);
593        assert_eq!(a.nrows(), 3);
594        assert_eq!(a.ncols(), 3);
595        assert_eq!(a.get_index(0, 0), f::<M>(2.0));
596        assert_eq!(a.get_index(1, 1), f::<M>(3.0));
597        assert_eq!(a.get_index(2, 2), f::<M>(5.0));
598        assert_eq!(a.get_index(0, 1), f::<M>(0.0));
599        assert_eq!(a.get_index(1, 0), f::<M>(0.0));
600    }
601
602    pub fn test_gemv<M: Matrix>() {
603        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
604        let values = vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)];
605        let a = M::try_from_triplets(2, 2, indices, values, Default::default()).unwrap();
606        let x = M::V::from_vec(vec![f::<M>(1.0), f::<M>(2.0)], Default::default());
607        let mut y = M::V::zeros(2, Default::default());
608        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
609        assert_eq!(y.clone_as_vec(), vec![f::<M>(5.0), f::<M>(11.0)]);
610
611        // beta = 1.0: accumulate onto the pre-existing value of y
612        a.gemv(f::<M>(1.0), &x, f::<M>(1.0), &mut y);
613        assert_eq!(y.clone_as_vec(), vec![f::<M>(10.0), f::<M>(22.0)]);
614    }
615
616    pub fn test_set_column<M: Matrix>() {
617        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
618        let values = vec![f::<M>(0.0), f::<M>(0.0), f::<M>(0.0), f::<M>(0.0)];
619        let mut a = M::try_from_triplets(2, 2, indices, values, Default::default()).unwrap();
620        let v = M::V::from_vec(vec![f::<M>(7.0), f::<M>(8.0)], Default::default());
621        a.set_column(1, &v);
622        let idx = triplet_indices(&a);
623        let vals = triplet_values(&a);
624        assert_eq!(idx, vec![(0, 0), (1, 0), (0, 1), (1, 1)]);
625        assert_eq!(
626            vals,
627            vec![f::<M>(0.0), f::<M>(0.0), f::<M>(7.0), f::<M>(8.0)]
628        );
629    }
630
631    pub fn test_copy_from<M: Matrix>() {
632        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
633        let values = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
634        let a = M::try_from_triplets(2, 2, indices, values, Default::default()).unwrap();
635        let mut b = M::zeros(2, 2, Default::default());
636        b.copy_from(&a);
637        let vals = triplet_values(&b);
638        assert_eq!(
639            vals,
640            vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)]
641        );
642    }
643
644    pub fn test_scale_add_and_assign<M: Matrix>() {
645        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
646        let x_vals = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
647        let y_vals = vec![f::<M>(10.0), f::<M>(20.0), f::<M>(30.0), f::<M>(40.0)];
648        let x = M::try_from_triplets(2, 2, indices.clone(), x_vals, Default::default()).unwrap();
649        let y = M::try_from_triplets(2, 2, indices, y_vals, Default::default()).unwrap();
650        let mut result = M::zeros(2, 2, Default::default());
651        result.copy_from(&x);
652        result.scale_add_and_assign(&x, f::<M>(2.0), &y);
653        let vals = triplet_values(&result);
654        assert_eq!(
655            vals,
656            vec![f::<M>(21.0), f::<M>(42.0), f::<M>(63.0), f::<M>(84.0)]
657        );
658    }
659
660    // --- DenseMatrix-specific tests ---
661
662    /// `Dense<M>`'s `try_from_indices`/`union` error branches: dense matrices ignore the actual
663    /// index pattern (their `Sparsity` is just a bounding box), so only a zero dimension or a
664    /// shape mismatch can make these fail — inputs the sparse backend wouldn't reject the same way.
665    pub fn test_sparsity_dense_errors<M: DenseMatrix>() {
666        assert!(M::Sparsity::try_from_indices(0, 3, vec![]).is_err());
667        let a = M::Sparsity::try_from_indices(2, 2, vec![]).unwrap();
668        let b = M::Sparsity::try_from_indices(3, 3, vec![]).unwrap();
669        assert!(a.union(b.as_ref()).is_err());
670    }
671
672    /// Order 1: 4 columns (`0..=order+2`), column 3 starts as garbage (overwritten).
673    /// Expected values are the right-to-left cumulative scan replayed by hand:
674    ///   col3 = d - col2
675    ///   col2 += col3 (updated); col1 += col2 (updated); col0 += col1 (updated)
676    pub fn test_update_backward_diff<M: DenseMatrix>() {
677        let order = 1;
678        let mut a = M::from_vec(
679            2,
680            4,
681            vec![
682                f::<M>(1.0),
683                f::<M>(10.0),
684                f::<M>(2.0),
685                f::<M>(20.0),
686                f::<M>(3.0),
687                f::<M>(30.0),
688                f::<M>(0.0),
689                f::<M>(0.0),
690            ],
691            Default::default(),
692        );
693        let d = M::V::from_vec(vec![f::<M>(100.0), f::<M>(1000.0)], Default::default());
694
695        a.update_backward_diff(order, &d);
696
697        assert_eq!(a.get_index(0, 3), f::<M>(97.0));
698        assert_eq!(a.get_index(1, 3), f::<M>(970.0));
699        assert_eq!(a.get_index(0, 2), f::<M>(100.0));
700        assert_eq!(a.get_index(1, 2), f::<M>(1000.0));
701        assert_eq!(a.get_index(0, 1), f::<M>(102.0));
702        assert_eq!(a.get_index(1, 1), f::<M>(1020.0));
703        assert_eq!(a.get_index(0, 0), f::<M>(103.0));
704        assert_eq!(a.get_index(1, 0), f::<M>(1030.0));
705    }
706
707    pub fn test_gemv_cols<M: DenseMatrix>() {
708        // columns: [1, 10], [2, 20], [3, 30], [4, 40]
709        let a = M::from_vec(
710            2,
711            4,
712            (1..=4)
713                .flat_map(|j| [f::<M>(j as f64), f::<M>(10.0 * j as f64)])
714                .collect(),
715            Default::default(),
716        );
717        let ones = |n: usize| vec![f::<M>(1.0); n];
718        let mut y = M::V::from_vec(vec![f::<M>(-1.0), f::<M>(-1.0)], Default::default());
719
720        // unweighted sum of columns 0..3, overwriting y
721        a.gemv_cols(0, 3, f::<M>(1.0), &ones(3), f::<M>(0.0), &mut y);
722        assert_eq!(y.clone_as_vec(), vec![f::<M>(6.0), f::<M>(60.0)]);
723
724        // weighted sum of every column
725        let w4 = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
726        a.gemv_cols(0, 4, f::<M>(1.0), &w4, f::<M>(0.0), &mut y);
727        assert_eq!(y.clone_as_vec(), vec![f::<M>(30.0), f::<M>(300.0)]);
728
729        // x lines up with the start of the range, not with column zero
730        let w23 = vec![f::<M>(2.0), f::<M>(3.0)];
731        a.gemv_cols(2, 4, f::<M>(1.0), &w23, f::<M>(0.0), &mut y);
732        assert_eq!(y.clone_as_vec(), vec![f::<M>(18.0), f::<M>(180.0)]);
733
734        // alpha scales the whole product
735        a.gemv_cols(2, 4, f::<M>(2.0), &w23, f::<M>(0.0), &mut y);
736        assert_eq!(y.clone_as_vec(), vec![f::<M>(36.0), f::<M>(360.0)]);
737
738        // beta = 1 accumulates onto y instead of overwriting it
739        let half = vec![f::<M>(0.5)];
740        a.gemv_cols(1, 2, f::<M>(1.0), &half, f::<M>(0.0), &mut y);
741        assert_eq!(y.clone_as_vec(), vec![f::<M>(1.0), f::<M>(10.0)]);
742        a.gemv_cols(1, 2, f::<M>(1.0), &half, f::<M>(1.0), &mut y);
743        assert_eq!(y.clone_as_vec(), vec![f::<M>(2.0), f::<M>(20.0)]);
744
745        // a general beta scales the old y before accumulating
746        a.gemv_cols(1, 2, f::<M>(1.0), &half, f::<M>(3.0), &mut y);
747        assert_eq!(y.clone_as_vec(), vec![f::<M>(7.0), f::<M>(70.0)]);
748
749        // an empty range contributes nothing, leaving y = beta * y
750        a.gemv_cols(1, 1, f::<M>(1.0), &ones(0), f::<M>(2.0), &mut y);
751        assert_eq!(y.clone_as_vec(), vec![f::<M>(14.0), f::<M>(140.0)]);
752        a.gemv_cols(1, 1, f::<M>(1.0), &ones(0), f::<M>(1.0), &mut y);
753        assert_eq!(y.clone_as_vec(), vec![f::<M>(14.0), f::<M>(140.0)]);
754        a.gemv_cols(1, 1, f::<M>(1.0), &ones(0), f::<M>(0.0), &mut y);
755        assert_eq!(y.clone_as_vec(), vec![f::<M>(0.0), f::<M>(0.0)]);
756    }
757
758    /// A column range wider than the 8-weight chunk the retired `weighted_column_sum` CUDA
759    /// kernel could take in one launch — that boundary was never covered by a test.
760    /// Entry `(i, j)` of batch `b`. `Matrix::get_index` only ever reads batch 0, and a
761    /// column view carries every batch, so this goes through the column.
762    fn mat_at<M: DenseMatrix>(m: &M, b: usize, i: usize, j: usize) -> M::T {
763        m.column(j).into_owned().clone_as_vec()[b * m.nrows() + i]
764    }
765
766    /// Every entry of batch `b`, column-major.
767    fn batch_as_vec<M: DenseMatrix>(m: &M, b: usize) -> Vec<M::T> {
768        let mut out = Vec::with_capacity(m.nrows() * m.ncols());
769        for j in 0..m.ncols() {
770            for i in 0..m.nrows() {
771                out.push(mat_at(m, b, i, j));
772            }
773        }
774        out
775    }
776
777    pub fn test_mul_cols_by<M: DenseMatrix>() {
778        // 2x4, columns [1,10], [2,20], [3,30], [4,40]
779        let mk = || {
780            M::from_vec(
781                2,
782                4,
783                (1..=4)
784                    .flat_map(|j| [f::<M>(j as f64), f::<M>(10.0 * j as f64)])
785                    .collect(),
786                Default::default(),
787            )
788        };
789        let untouched = batch_as_vec(&mk(), 0);
790
791        // identity must leave the matrix exactly as it was: the sharpest check that the
792        // in-place update never reads a value it has already overwritten
793        let mut a = mk();
794        a.mul_cols_by(2, &[f::<M>(1.0), f::<M>(0.0), f::<M>(0.0), f::<M>(1.0)]);
795        assert_eq!(batch_as_vec(&a, 0), untouched);
796
797        // an empty range is a no-op
798        let mut a = mk();
799        a.mul_cols_by(0, &[]);
800        assert_eq!(batch_as_vec(&a, 0), untouched);
801
802        // a single column is just scaled, and the rest are left alone
803        let mut a = mk();
804        a.mul_cols_by(1, &[f::<M>(3.0)]);
805        assert_eq!(
806            batch_as_vec(&a, 0),
807            vec![
808                f::<M>(3.0),
809                f::<M>(30.0),
810                f::<M>(2.0),
811                f::<M>(20.0),
812                f::<M>(3.0),
813                f::<M>(30.0),
814                f::<M>(4.0),
815                f::<M>(40.0)
816            ]
817        );
818
819        // rhs = [[1, 2], [3, 4]] column-major, applied to the first two columns only:
820        //   new[:, 0] = 1 * [1,10] + 3 * [2,20] = [7, 70]
821        //   new[:, 1] = 2 * [1,10] + 4 * [2,20] = [10, 100]
822        // columns 2 and 3 must survive untouched -- the BDF difference table keeps live data
823        // above the updated range and reads it back in `update_backward_diff`
824        let mut a = mk();
825        a.mul_cols_by(2, &[f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)]);
826        assert_eq!(
827            batch_as_vec(&a, 0),
828            vec![
829                f::<M>(7.0),
830                f::<M>(70.0),
831                f::<M>(10.0),
832                f::<M>(100.0),
833                f::<M>(3.0),
834                f::<M>(30.0),
835                f::<M>(4.0),
836                f::<M>(40.0)
837            ]
838        );
839
840        // the full width the solvers use (MAX_ORDER + 1), via an identity
841        let ncols = 6;
842        let mut wide = M::from_vec(
843            2,
844            ncols,
845            (1..=ncols)
846                .flat_map(|j| [f::<M>(j as f64), f::<M>(10.0 * j as f64)])
847                .collect(),
848            Default::default(),
849        );
850        let wide_before = batch_as_vec(&wide, 0);
851        let mut id = vec![f::<M>(0.0); ncols * ncols];
852        for j in 0..ncols {
853            id[j * ncols + j] = f::<M>(1.0);
854        }
855        wide.mul_cols_by(ncols, &id);
856        assert_eq!(batch_as_vec(&wide, 0), wide_before);
857
858        // more rows than fit one staging tile, so the row-blocked loop runs several times
859        let nrows = 700;
860        let mut tall = M::from_vec(
861            nrows,
862            2,
863            (0..2 * nrows).map(|i| f::<M>(i as f64)).collect(),
864            Default::default(),
865        );
866        tall.mul_cols_by(2, &[f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)]);
867        let got = batch_as_vec(&tall, 0);
868        for i in 0..nrows {
869            let (c0, c1) = (f::<M>(i as f64), f::<M>((nrows + i) as f64));
870            assert_eq!(got[i], c0 + f::<M>(3.0) * c1);
871            assert_eq!(got[nrows + i], f::<M>(2.0) * c0 + f::<M>(4.0) * c1);
872        }
873    }
874
875    /// Same product as [`test_mul_cols_by`] on 2 independent batches: `rhs` is shared by every
876    /// batch, the columns are not.
877    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
878    pub fn test_batched_mul_cols_by<M: DenseMatrix>(ctx: M::C) {
879        assert_eq!(ctx.nbatch(), 2);
880        // batch0 columns [1,10], [2,20]; batch1 columns [3,30], [4,40]
881        let mut a = M::from_vec(
882            2,
883            2,
884            (1..=4)
885                .flat_map(|j| [f::<M>(j as f64), f::<M>(10.0 * j as f64)])
886                .collect(),
887            ctx,
888        );
889        a.mul_cols_by(2, &[f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)]);
890        // batch0: [1,10] + 3*[2,20] = [7,70];   2*[1,10] + 4*[2,20] = [10,100]
891        assert_eq!(
892            batch_as_vec(&a, 0),
893            vec![f::<M>(7.0), f::<M>(70.0), f::<M>(10.0), f::<M>(100.0)]
894        );
895        // batch1: [3,30] + 3*[4,40] = [15,150]; 2*[3,30] + 4*[4,40] = [22,220]
896        assert_eq!(
897            batch_as_vec(&a, 1),
898            vec![f::<M>(15.0), f::<M>(150.0), f::<M>(22.0), f::<M>(220.0)]
899        );
900    }
901
902    /// The widest column range `gemv_cols` accepts, which is also the size of the by-value
903    /// coefficient array the CUDA kernel takes.
904    pub fn test_gemv_cols_many_columns<M: DenseMatrix>() {
905        let ncols = MAX_SMALL_COLS;
906        // column j is [j + 1, 10 * (j + 1)]
907        let a = M::from_vec(
908            2,
909            ncols,
910            (1..=ncols)
911                .flat_map(|j| [f::<M>(j as f64), f::<M>(10.0 * j as f64)])
912                .collect(),
913            Default::default(),
914        );
915        let x = vec![f::<M>(1.0); ncols];
916        let mut y = M::V::from_vec(vec![f::<M>(-1.0), f::<M>(-1.0)], Default::default());
917        a.gemv_cols(0, ncols, f::<M>(1.0), &x, f::<M>(0.0), &mut y);
918        // sum_{j=1}^{ncols} j
919        let total = f::<M>((ncols * (ncols + 1) / 2) as f64);
920        assert_eq!(y.clone_as_vec(), vec![total, f::<M>(10.0) * total]);
921    }
922
923    /// Same products as [`test_gemv_cols`], run on 2 independent batches: `x` is unbatched and
924    /// therefore shared by every batch, the columns are not.
925    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
926    pub fn test_batched_gemv_cols_m<M: DenseMatrix>(ctx: M::C) {
927        assert_eq!(ctx.nbatch(), 2);
928
929        // batch 0 columns: [1, 10], [2, 20]; batch 1 columns: [3, 30], [4, 40]
930        let a = M::from_vec(
931            2,
932            2,
933            (1..=4)
934                .flat_map(|j| [f::<M>(j as f64), f::<M>(10.0 * j as f64)])
935                .collect(),
936            ctx.clone(),
937        );
938        let mut y = M::V::from_vec(
939            vec![f::<M>(-1.0), f::<M>(-1.0), f::<M>(-1.0), f::<M>(-1.0)],
940            ctx,
941        );
942
943        let ones = vec![f::<M>(1.0), f::<M>(1.0)];
944        a.gemv_cols(0, 2, f::<M>(1.0), &ones, f::<M>(0.0), &mut y);
945        assert_eq!(
946            y.clone_as_vec(),
947            vec![f::<M>(3.0), f::<M>(30.0), f::<M>(7.0), f::<M>(70.0)]
948        );
949
950        let w23 = vec![f::<M>(2.0), f::<M>(3.0)];
951        a.gemv_cols(0, 2, f::<M>(1.0), &w23, f::<M>(0.0), &mut y);
952        assert_eq!(
953            y.clone_as_vec(),
954            vec![f::<M>(8.0), f::<M>(80.0), f::<M>(18.0), f::<M>(180.0)]
955        );
956
957        let half = vec![f::<M>(0.5)];
958        a.gemv_cols(1, 2, f::<M>(1.0), &half, f::<M>(0.0), &mut y);
959        assert_eq!(
960            y.clone_as_vec(),
961            vec![f::<M>(1.0), f::<M>(10.0), f::<M>(2.0), f::<M>(20.0)]
962        );
963
964        // beta = 1 accumulates independently in each batch
965        a.gemv_cols(1, 2, f::<M>(1.0), &half, f::<M>(1.0), &mut y);
966        assert_eq!(
967            y.clone_as_vec(),
968            vec![f::<M>(2.0), f::<M>(20.0), f::<M>(4.0), f::<M>(40.0)]
969        );
970    }
971
972    pub fn test_resize_cols<M: DenseMatrix>() {
973        let mut a = M::zeros(2, 2, Default::default());
974        a.set_index(0, 0, M::T::one());
975        a.set_index(0, 1, M::T::from_f64(2.0).unwrap());
976        a.set_index(1, 0, M::T::from_f64(3.0).unwrap());
977        a.set_index(1, 1, M::T::from_f64(4.0).unwrap());
978
979        // resizing to the current column count is a no-op
980        a.resize_cols(2);
981        assert_eq!(a.ncols(), 2);
982        assert_eq!(a.get_index(0, 0), M::T::one());
983
984        a.resize_cols(3);
985        assert_eq!(a.ncols(), 3);
986        assert_eq!(a.nrows(), 2);
987        assert_eq!(a.get_index(0, 0), M::T::one());
988        assert_eq!(a.get_index(0, 1), M::T::from_f64(2.0).unwrap());
989        assert_eq!(a.get_index(1, 0), M::T::from_f64(3.0).unwrap());
990        assert_eq!(a.get_index(1, 1), M::T::from_f64(4.0).unwrap());
991
992        a.set_index(0, 2, M::T::from_f64(5.0).unwrap());
993        a.set_index(1, 2, M::T::from_f64(6.0).unwrap());
994        assert_eq!(a.get_index(0, 2), M::T::from_f64(5.0).unwrap());
995        assert_eq!(a.get_index(1, 2), M::T::from_f64(6.0).unwrap());
996
997        a.resize_cols(2);
998        assert_eq!(a.ncols(), 2);
999        assert_eq!(a.nrows(), 2);
1000        assert_eq!(a.get_index(0, 0), M::T::one());
1001        assert_eq!(a.get_index(0, 1), M::T::from_f64(2.0).unwrap());
1002        assert_eq!(a.get_index(1, 0), M::T::from_f64(3.0).unwrap());
1003        assert_eq!(a.get_index(1, 1), M::T::from_f64(4.0).unwrap());
1004    }
1005
1006    pub fn test_from_vec<M: DenseMatrix>() {
1007        let a = M::from_vec(
1008            2,
1009            2,
1010            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
1011            Default::default(),
1012        );
1013        assert_eq!(a.nrows(), 2);
1014        assert_eq!(a.ncols(), 2);
1015        assert_eq!(a.get_index(0, 0), f::<M>(1.0));
1016        assert_eq!(a.get_index(1, 0), f::<M>(3.0));
1017        assert_eq!(a.get_index(0, 1), f::<M>(2.0));
1018        assert_eq!(a.get_index(1, 1), f::<M>(4.0));
1019    }
1020
1021    pub fn test_column_view<M: DenseMatrix>() {
1022        let a = M::from_vec(
1023            2,
1024            2,
1025            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
1026            Default::default(),
1027        );
1028        let col = a.column(1);
1029        use crate::VectorView;
1030        assert_eq!(col.get_index(0), f::<M>(2.0));
1031        assert_eq!(col.get_index(1), f::<M>(4.0));
1032    }
1033
1034    // --- Batched Matrix-generic tests ---
1035
1036    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1037    pub fn test_batched_zeros_m<M: Matrix>(ctx: M::C) {
1038        assert_eq!(ctx.nbatch(), 2);
1039        let a = M::zeros(2, 3, ctx);
1040        assert_eq!(a.nrows(), 2);
1041        assert_eq!(a.ncols(), 3);
1042        let vals = triplet_values(&a);
1043        assert!(vals.is_empty() || vals.iter().all(|v| v.is_zero()));
1044    }
1045
1046    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1047    pub fn test_batched_gemv_m<M: Matrix>(ctx: M::C) {
1048        assert_eq!(ctx.nbatch(), 2);
1049        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1050        let values = vec![
1051            f::<M>(1.0),
1052            f::<M>(3.0),
1053            f::<M>(2.0),
1054            f::<M>(4.0), // batch 0
1055            f::<M>(5.0),
1056            f::<M>(7.0),
1057            f::<M>(6.0),
1058            f::<M>(8.0), // batch 1
1059        ];
1060        let a = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
1061        let x = M::V::from_vec(
1062            vec![f::<M>(1.0), f::<M>(2.0), f::<M>(1.0), f::<M>(1.0)],
1063            ctx.clone(),
1064        );
1065        let mut y = M::V::zeros(2, ctx);
1066        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1067        assert_eq!(
1068            y.clone_as_vec(),
1069            vec![f::<M>(5.0), f::<M>(11.0), f::<M>(11.0), f::<M>(15.0)]
1070        );
1071    }
1072
1073    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1074    pub fn test_batched_gemv_broadcast_x_m<M: Matrix>(ctx: M::C) {
1075        assert_eq!(ctx.nbatch(), 2);
1076        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1077        let values = vec![
1078            f::<M>(1.0),
1079            f::<M>(3.0),
1080            f::<M>(2.0),
1081            f::<M>(4.0),
1082            f::<M>(5.0),
1083            f::<M>(7.0),
1084            f::<M>(6.0),
1085            f::<M>(8.0),
1086        ];
1087        let a = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
1088        let x = M::V::from_vec(vec![f::<M>(1.0), f::<M>(2.0)], Default::default());
1089        let mut y = M::V::zeros(2, ctx);
1090        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1091        assert_eq!(
1092            y.clone_as_vec(),
1093            vec![f::<M>(5.0), f::<M>(11.0), f::<M>(17.0), f::<M>(23.0)]
1094        );
1095    }
1096
1097    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1098    pub fn test_batched_gemv_broadcast_mat_m<M: Matrix>(ctx: M::C) {
1099        assert_eq!(ctx.nbatch(), 2);
1100        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1101        let values = vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)];
1102        let a =
1103            M::try_from_triplets(2, 2, indices, values, ctx.clone_with_nbatch(1).unwrap()).unwrap();
1104        let x = M::V::from_vec(
1105            vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)],
1106            ctx.clone(),
1107        );
1108        let mut y = M::V::zeros(2, ctx);
1109        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1110        assert_eq!(
1111            y.clone_as_vec(),
1112            vec![f::<M>(5.0), f::<M>(11.0), f::<M>(11.0), f::<M>(25.0)]
1113        );
1114    }
1115
1116    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1117    pub fn test_batched_from_diagonal_m<M: Matrix>(ctx: M::C) {
1118        assert_eq!(ctx.nbatch(), 2);
1119        let v = M::V::from_vec(
1120            vec![f::<M>(2.0), f::<M>(3.0), f::<M>(4.0), f::<M>(5.0)],
1121            ctx,
1122        );
1123        let a = M::from_diagonal(&v);
1124        assert_eq!(a.nrows(), 2);
1125        assert_eq!(a.ncols(), 2);
1126        let idx = triplet_indices(&a);
1127        let vals = triplet_values(&a);
1128        for &(i, j) in &idx {
1129            let pos = idx.iter().position(|&x| x == (i, j)).unwrap();
1130            if i == j {
1131                assert!(
1132                    vals[pos] != M::T::zero(),
1133                    "diagonal entry should be non-zero"
1134                );
1135            } else {
1136                assert!(vals[pos].is_zero(), "off-diagonal entry should be zero");
1137            }
1138        }
1139    }
1140
1141    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1142    pub fn test_batched_copy_from_m<M: Matrix>(ctx: M::C) {
1143        assert_eq!(ctx.nbatch(), 2);
1144        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1145        let values = vec![
1146            f::<M>(1.0),
1147            f::<M>(2.0),
1148            f::<M>(3.0),
1149            f::<M>(4.0),
1150            f::<M>(5.0),
1151            f::<M>(6.0),
1152            f::<M>(7.0),
1153            f::<M>(8.0),
1154        ];
1155        let a = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
1156        let mut b = M::zeros(2, 2, ctx);
1157        b.copy_from(&a);
1158        let vals = triplet_values(&b);
1159        assert_eq!(
1160            vals,
1161            vec![
1162                f::<M>(1.0),
1163                f::<M>(2.0),
1164                f::<M>(3.0),
1165                f::<M>(4.0),
1166                f::<M>(5.0),
1167                f::<M>(6.0),
1168                f::<M>(7.0),
1169                f::<M>(8.0),
1170            ]
1171        );
1172    }
1173
1174    /// `copy_from` with mismatched batch counts: the source broadcasts over the destination's
1175    /// batches, exercising the per-batch loop instead of the equal-nbatch whole-matrix copy.
1176    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1177    pub fn test_batched_copy_from_broadcast_m<M: Matrix>(ctx: M::C) {
1178        assert_eq!(ctx.nbatch(), 2);
1179        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1180        let values = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
1181        let a = M::try_from_triplets(2, 2, indices, values, M::C::default()).unwrap();
1182        let mut b = M::zeros(2, 2, ctx);
1183        b.copy_from(&a);
1184        let vals = triplet_values(&b);
1185        assert_eq!(
1186            vals,
1187            vec![
1188                f::<M>(1.0),
1189                f::<M>(2.0),
1190                f::<M>(3.0),
1191                f::<M>(4.0),
1192                f::<M>(1.0),
1193                f::<M>(2.0),
1194                f::<M>(3.0),
1195                f::<M>(4.0),
1196            ]
1197        );
1198    }
1199
1200    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1201    pub fn test_batched_set_column_m<M: Matrix>(ctx: M::C) {
1202        assert_eq!(ctx.nbatch(), 2);
1203        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1204        let values = vec![
1205            f::<M>(0.0),
1206            f::<M>(0.0),
1207            f::<M>(0.0),
1208            f::<M>(0.0),
1209            f::<M>(0.0),
1210            f::<M>(0.0),
1211            f::<M>(0.0),
1212            f::<M>(0.0),
1213        ];
1214        let mut a = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
1215        let v = M::V::from_vec(
1216            vec![f::<M>(5.0), f::<M>(6.0), f::<M>(7.0), f::<M>(8.0)],
1217            ctx,
1218        );
1219        a.set_column(0, &v);
1220        let vals = triplet_values(&a);
1221        assert_eq!(
1222            vals,
1223            vec![
1224                f::<M>(5.0),
1225                f::<M>(6.0),
1226                f::<M>(0.0),
1227                f::<M>(0.0),
1228                f::<M>(7.0),
1229                f::<M>(8.0),
1230                f::<M>(0.0),
1231                f::<M>(0.0),
1232            ]
1233        );
1234    }
1235
1236    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1237    pub fn test_batched_scale_add_and_assign_m<M: Matrix>(ctx: M::C) {
1238        assert_eq!(ctx.nbatch(), 2);
1239        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1240        let x_vals = vec![
1241            f::<M>(1.0),
1242            f::<M>(2.0),
1243            f::<M>(3.0),
1244            f::<M>(4.0),
1245            f::<M>(5.0),
1246            f::<M>(6.0),
1247            f::<M>(7.0),
1248            f::<M>(8.0),
1249        ];
1250        let y_vals = vec![
1251            f::<M>(10.0),
1252            f::<M>(20.0),
1253            f::<M>(30.0),
1254            f::<M>(40.0),
1255            f::<M>(50.0),
1256            f::<M>(60.0),
1257            f::<M>(70.0),
1258            f::<M>(80.0),
1259        ];
1260        let x = M::try_from_triplets(2, 2, indices.clone(), x_vals, ctx.clone()).unwrap();
1261        let y = M::try_from_triplets(2, 2, indices, y_vals, ctx.clone()).unwrap();
1262        let mut result = M::zeros(2, 2, ctx);
1263        result.copy_from(&x);
1264        result.scale_add_and_assign(&x, f::<M>(2.0), &y);
1265        let vals = triplet_values(&result);
1266        assert_eq!(
1267            vals,
1268            vec![
1269                f::<M>(21.0),
1270                f::<M>(42.0),
1271                f::<M>(63.0),
1272                f::<M>(84.0),
1273                f::<M>(105.0),
1274                f::<M>(126.0),
1275                f::<M>(147.0),
1276                f::<M>(168.0),
1277            ]
1278        );
1279    }
1280
1281    /// An owned right-hand side may be written into, which reverses the operands of `sub`.
1282    ///
1283    /// Only wired up for the backends that implement the owned-rhs combinations; move the call
1284    /// into `generate_dense_matrix_tests_nonbatched!` once they all do.
1285    pub fn test_owned_rhs_m<M>()
1286    where
1287        M: DenseMatrix,
1288        for<'a> &'a M: Add<M, Output = M> + Sub<M, Output = M>,
1289    {
1290        let values = vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)];
1291        let a = M::from_vec(
1292            2,
1293            2,
1294            vec![f::<M>(10.0), f::<M>(30.0), f::<M>(20.0), f::<M>(40.0)],
1295            Default::default(),
1296        );
1297
1298        let b = M::from_vec(2, 2, values.clone(), Default::default());
1299        let c = &a - b;
1300        assert_eq!(c.get_index(0, 0), f::<M>(9.0));
1301        assert_eq!(c.get_index(1, 1), f::<M>(36.0));
1302
1303        let b = M::from_vec(2, 2, values, Default::default());
1304        let c = &a + b;
1305        assert_eq!(c.get_index(0, 0), f::<M>(11.0));
1306        assert_eq!(c.get_index(1, 1), f::<M>(44.0));
1307    }
1308
1309    /// Owned right-hand side with broadcasting in both directions: the right-hand side can
1310    /// only be written into when it already has the result's batch count.
1311    pub fn test_batched_owned_rhs_broadcast_m<M>(ctx: M::C)
1312    where
1313        M: Matrix + DenseMatrix,
1314        for<'a> &'a M: Sub<M, Output = M>,
1315    {
1316        assert_eq!(ctx.nbatch(), 2);
1317        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1318        let one_batch = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
1319        let two_batches = vec![
1320            f::<M>(10.0),
1321            f::<M>(20.0),
1322            f::<M>(30.0),
1323            f::<M>(40.0),
1324            f::<M>(50.0),
1325            f::<M>(60.0),
1326            f::<M>(70.0),
1327            f::<M>(80.0),
1328        ];
1329
1330        // lhs broadcasts over the batches of rhs, which already spans the result
1331        let a = M::try_from_triplets(2, 2, indices.clone(), one_batch.clone(), M::C::default())
1332            .unwrap();
1333        let b =
1334            M::try_from_triplets(2, 2, indices.clone(), two_batches.clone(), ctx.clone()).unwrap();
1335        let c = &a - b;
1336        assert_eq!(c.context().nbatch(), 2);
1337        assert_eq!(
1338            triplet_values(&c),
1339            vec![
1340                f::<M>(-9.0),
1341                f::<M>(-18.0),
1342                f::<M>(-27.0),
1343                f::<M>(-36.0),
1344                f::<M>(-49.0),
1345                f::<M>(-58.0),
1346                f::<M>(-67.0),
1347                f::<M>(-76.0),
1348            ]
1349        );
1350    }
1351
1352    /// The owned right-hand side governs, so a *wider* borrowed left-hand side is an error.
1353    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1354    pub fn test_batched_owned_rhs_narrow_m<M>(ctx: M::C)
1355    where
1356        M: Matrix + DenseMatrix,
1357        for<'a> &'a M: Sub<M, Output = M>,
1358    {
1359        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1360        let wide = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1361        let narrow = (1..=4).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1362        let a = M::try_from_triplets(2, 2, indices.clone(), wide, ctx).unwrap();
1363        let b = M::try_from_triplets(2, 2, indices, narrow, M::C::default()).unwrap();
1364        let _ = &a - b;
1365    }
1366
1367    /// The owned left-hand side governs and is written into.
1368    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1369    pub fn test_batched_owned_lhs_broadcast_m<M>(ctx: M::C)
1370    where
1371        M: Matrix + for<'a> Sub<&'a M, Output = M>,
1372    {
1373        assert_eq!(ctx.nbatch(), 2);
1374        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1375        let one_batch = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
1376        let two_batches = vec![
1377            f::<M>(10.0),
1378            f::<M>(20.0),
1379            f::<M>(30.0),
1380            f::<M>(40.0),
1381            f::<M>(50.0),
1382            f::<M>(60.0),
1383            f::<M>(70.0),
1384            f::<M>(80.0),
1385        ];
1386
1387        // rhs broadcasts over the batches of lhs
1388        let a =
1389            M::try_from_triplets(2, 2, indices.clone(), two_batches.clone(), ctx.clone()).unwrap();
1390        let b = M::try_from_triplets(2, 2, indices.clone(), one_batch.clone(), M::C::default())
1391            .unwrap();
1392        let c = a - &b;
1393        assert_eq!(c.context().nbatch(), 2);
1394        assert_eq!(
1395            triplet_values(&c),
1396            vec![
1397                f::<M>(9.0),
1398                f::<M>(18.0),
1399                f::<M>(27.0),
1400                f::<M>(36.0),
1401                f::<M>(49.0),
1402                f::<M>(58.0),
1403                f::<M>(67.0),
1404                f::<M>(76.0),
1405            ]
1406        );
1407    }
1408
1409    /// The owned left-hand side governs, so a *wider* right-hand side is an error.
1410    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1411    pub fn test_batched_owned_lhs_narrow_m<M>(ctx: M::C)
1412    where
1413        M: Matrix + for<'a> Sub<&'a M, Output = M>,
1414    {
1415        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1416        let narrow = (1..=4).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1417        let wide = (1..=8).map(|i| f::<M>(10.0 * i as f64)).collect::<Vec<_>>();
1418        let a = M::try_from_triplets(2, 2, indices.clone(), narrow, M::C::default()).unwrap();
1419        let b = M::try_from_triplets(2, 2, indices, wide, ctx).unwrap();
1420        let _ = a - &b;
1421    }
1422
1423    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1424    pub fn test_batched_scale_add_and_assign_broadcast_m<M: Matrix>(ctx: M::C) {
1425        assert_eq!(ctx.nbatch(), 2);
1426        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1427        let x_vals = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
1428        let y_vals = vec![f::<M>(10.0), f::<M>(20.0), f::<M>(30.0), f::<M>(40.0)];
1429        // x and y have nbatch == 1, so they broadcast over both batches of result
1430        let x =
1431            M::try_from_triplets(2, 2, indices.clone(), x_vals.clone(), M::C::default()).unwrap();
1432        let y = M::try_from_triplets(2, 2, indices.clone(), y_vals, M::C::default()).unwrap();
1433        // result carries the union sparsity (required for sparse matrices) in both batches
1434        let mut result_vals = x_vals.clone();
1435        result_vals.extend(x_vals);
1436        let mut result = M::try_from_triplets(2, 2, indices, result_vals, ctx).unwrap();
1437        result.scale_add_and_assign(&x, f::<M>(2.0), &y);
1438        let vals = triplet_values(&result);
1439        assert_eq!(
1440            vals,
1441            vec![
1442                f::<M>(21.0),
1443                f::<M>(42.0),
1444                f::<M>(63.0),
1445                f::<M>(84.0),
1446                f::<M>(21.0),
1447                f::<M>(42.0),
1448                f::<M>(63.0),
1449                f::<M>(84.0),
1450            ]
1451        );
1452    }
1453
1454    /// Scaling through a reference is a separate `impl` from the owned one, and it walks the
1455    /// batches itself.
1456    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1457    pub fn test_batched_mul_scalar_ref_m<M>(ctx: M::C)
1458    where
1459        M: Matrix,
1460        for<'a> &'a M: Mul<Scale<M::T>, Output = M>,
1461    {
1462        assert_eq!(ctx.nbatch(), 2);
1463        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1464        let values = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1465        let a = M::try_from_triplets(2, 2, indices, values, ctx).unwrap();
1466        let result = &a * Scale(f::<M>(2.0));
1467        assert_eq!(result.context().nbatch(), 2);
1468        assert_eq!(
1469            triplet_values(&result),
1470            (1..=8).map(|i| f::<M>(2.0 * i as f64)).collect::<Vec<_>>()
1471        );
1472    }
1473
1474    /// `+=` with a narrower right-hand side: the equal-batch whole-matrix path cannot be
1475    /// used, so this is the only test of the per-batch assign loop.
1476    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1477    pub fn test_batched_add_assign_broadcast_m<M: DenseMatrix>(ctx: M::C) {
1478        assert_eq!(ctx.nbatch(), 2);
1479        let mut a = M::from_vec(
1480            2,
1481            2,
1482            (1..=8).map(|i| f::<M>(i as f64)).collect(),
1483            ctx.clone(),
1484        );
1485        let b = M::from_vec(
1486            2,
1487            2,
1488            (1..=4).map(|i| f::<M>(10.0 * i as f64)).collect(),
1489            M::C::default(),
1490        );
1491        a += &b;
1492        // batch 0: 1..4 + 10..40, batch 1: 5..8 + 10..40
1493        let expected = (1..=4)
1494            .map(|i| f::<M>(i as f64 + 10.0 * i as f64))
1495            .chain((5..=8).map(|i| f::<M>(i as f64 + 10.0 * (i - 4) as f64)))
1496            .collect::<Vec<_>>();
1497        assert_eq!(triplet_values(&a), expected);
1498    }
1499
1500    /// Owned left-hand side broadcast into a fresh allocation by `+` (the `-` direction is
1501    /// covered by [`test_batched_owned_lhs_broadcast_m`]).
1502    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1503    pub fn test_batched_owned_lhs_add_broadcast_m<M>(ctx: M::C)
1504    where
1505        M: Matrix + for<'a> Add<&'a M, Output = M>,
1506    {
1507        assert_eq!(ctx.nbatch(), 2);
1508        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1509        let narrow = M::try_from_triplets(
1510            2,
1511            2,
1512            indices.clone(),
1513            (1..=4).map(|i| f::<M>(i as f64)).collect(),
1514            M::C::default(),
1515        )
1516        .unwrap();
1517        let wide = M::try_from_triplets(
1518            2,
1519            2,
1520            indices,
1521            (1..=8).map(|i| f::<M>(10.0 * i as f64)).collect(),
1522            ctx,
1523        )
1524        .unwrap();
1525        // the owned left-hand side governs, so it is the wide one and `narrow` broadcasts in
1526        let c = wide + &narrow;
1527        assert_eq!(c.context().nbatch(), 2);
1528        let expected = (1..=4)
1529            .map(|i| f::<M>(10.0 * i as f64 + i as f64))
1530            .chain((5..=8).map(|i| f::<M>(10.0 * i as f64 + (i - 4) as f64)))
1531            .collect::<Vec<_>>();
1532        assert_eq!(triplet_values(&c), expected);
1533    }
1534
1535    /// `gemv` with `beta` neither 0 nor 1, which is the only path that rescales `y` before
1536    /// accumulating.
1537    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1538    pub fn test_batched_gemv_beta_m<M: Matrix>(ctx: M::C) {
1539        assert_eq!(ctx.nbatch(), 2);
1540        // batch 0 = I, batch 1 = 2I
1541        let indices = vec![(0, 0), (1, 1)];
1542        let values = vec![f::<M>(1.0), f::<M>(1.0), f::<M>(2.0), f::<M>(2.0)];
1543        let a = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
1544        let x = M::V::from_vec(
1545            vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)],
1546            ctx.clone(),
1547        );
1548        let mut y = M::V::from_vec(
1549            vec![f::<M>(1.0), f::<M>(1.0), f::<M>(1.0), f::<M>(1.0)],
1550            ctx,
1551        );
1552        // y = 1 * A * x + 2 * y
1553        a.gemv(f::<M>(1.0), &x, f::<M>(2.0), &mut y);
1554        assert_eq!(
1555            y.clone_as_vec(),
1556            vec![f::<M>(3.0), f::<M>(4.0), f::<M>(8.0), f::<M>(10.0)]
1557        );
1558    }
1559
1560    // --- Grouped broadcast tests: `B` batches feeding `B * P` batches ---
1561    //
1562    // Each group holds a distinct value, so a cyclic mapping (source order `[0, 1, 0, 1]`)
1563    // fails these — grouped broadcasting reads `[0, 0, 1, 1]`.
1564
1565    /// `nbatch = 2` widened to 4, so each source batch covers two contiguous destinations.
1566    fn ctx4<M: Matrix>(ctx2: &M::C) -> M::C {
1567        assert_eq!(ctx2.nbatch(), 2);
1568        ctx2.clone_with_nbatch(4).unwrap()
1569    }
1570
1571    /// `y` carries 4 batches while the matrix carries 2, so each factor of the matrix serves
1572    /// two contiguous right-hand sides.
1573    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1574    pub fn test_grouped_gemv_m<M: Matrix>(ctx2: M::C) {
1575        let wide = ctx4::<M>(&ctx2);
1576        // batch 0 = I, batch 1 = 2I
1577        let indices = vec![(0, 0), (1, 1)];
1578        let values = vec![f::<M>(1.0), f::<M>(1.0), f::<M>(2.0), f::<M>(2.0)];
1579        let a = M::try_from_triplets(2, 2, indices, values, ctx2).unwrap();
1580        let x = M::V::from_vec((1..=8).map(|i| f::<M>(i as f64)).collect(), wide.clone());
1581        let mut y = M::V::zeros(2, wide);
1582        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1583        assert_eq!(
1584            y.clone_as_vec(),
1585            vec![
1586                f::<M>(1.0),
1587                f::<M>(2.0),
1588                f::<M>(3.0),
1589                f::<M>(4.0),
1590                f::<M>(10.0),
1591                f::<M>(12.0),
1592                f::<M>(14.0),
1593                f::<M>(16.0),
1594            ]
1595        );
1596    }
1597
1598    /// The mirror of [`test_grouped_gemv_m`]: the matrix holds one lane per right-hand side
1599    /// while `x` carries 2, so each `x` lane serves two contiguous results — which still
1600    /// multiply different matrix lanes.
1601    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1602    pub fn test_grouped_gemv_x_m<M: Matrix>(ctx2: M::C) {
1603        let wide = ctx4::<M>(&ctx2);
1604        // batch b = (b + 1) I
1605        let indices = vec![(0, 0), (1, 1)];
1606        let values = (1..=4).flat_map(|i| [f::<M>(i as f64); 2]).collect();
1607        let a = M::try_from_triplets(2, 2, indices, values, wide.clone()).unwrap();
1608        let x = M::V::from_vec((1..=4).map(|i| f::<M>(i as f64)).collect(), ctx2);
1609        let mut y = M::V::zeros(2, wide);
1610        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1611        assert_eq!(
1612            y.clone_as_vec(),
1613            vec![
1614                f::<M>(1.0),
1615                f::<M>(2.0),
1616                f::<M>(2.0),
1617                f::<M>(4.0),
1618                f::<M>(9.0),
1619                f::<M>(12.0),
1620                f::<M>(12.0),
1621                f::<M>(16.0),
1622            ]
1623        );
1624    }
1625
1626    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1627    pub fn test_grouped_copy_from_m<M: Matrix>(ctx2: M::C) {
1628        let wide = ctx4::<M>(&ctx2);
1629        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1630        let values = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1631        let a = M::try_from_triplets(2, 2, indices, values, ctx2).unwrap();
1632        let mut b = M::zeros(2, 2, wide);
1633        b.copy_from(&a);
1634        let group0 = (1..=4).map(|i| f::<M>(i as f64));
1635        let group1 = (5..=8).map(|i| f::<M>(i as f64));
1636        let expected = group0
1637            .clone()
1638            .chain(group0)
1639            .chain(group1.clone())
1640            .chain(group1)
1641            .collect::<Vec<_>>();
1642        assert_eq!(triplet_values(&b), expected);
1643    }
1644
1645    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1646    pub fn test_grouped_scale_add_and_assign_m<M: Matrix>(ctx2: M::C) {
1647        let wide = ctx4::<M>(&ctx2);
1648        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1649        let x_vals = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1650        let y_vals = (1..=8).map(|i| f::<M>(10.0 * i as f64)).collect::<Vec<_>>();
1651        let x = M::try_from_triplets(2, 2, indices.clone(), x_vals, ctx2.clone()).unwrap();
1652        let y = M::try_from_triplets(2, 2, indices.clone(), y_vals, ctx2).unwrap();
1653        // the destination needs the union sparsity in every batch (required for sparse)
1654        let dst_vals = (0..16).map(|_| f::<M>(0.0)).collect::<Vec<_>>();
1655        let mut result = M::try_from_triplets(2, 2, indices, dst_vals, wide).unwrap();
1656        result.scale_add_and_assign(&x, f::<M>(2.0), &y);
1657        // x_i + 2 * y_i = 21 * i, with each source batch covering two contiguous batches
1658        let group0 = (1..=4).map(|i| f::<M>(21.0 * i as f64));
1659        let group1 = (5..=8).map(|i| f::<M>(21.0 * i as f64));
1660        let expected = group0
1661            .clone()
1662            .chain(group0)
1663            .chain(group1.clone())
1664            .chain(group1)
1665            .collect::<Vec<_>>();
1666        assert_eq!(triplet_values(&result), expected);
1667    }
1668
1669    /// The grouped flavour of [`test_batched_add_assign_broadcast_m`].
1670    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1671    pub fn test_grouped_add_assign_m<M: DenseMatrix>(ctx2: M::C) {
1672        let wide = ctx4::<M>(&ctx2);
1673        let mut a = M::from_vec(2, 2, (0..16).map(|_| f::<M>(0.0)).collect(), wide);
1674        let b = M::from_vec(2, 2, (1..=8).map(|i| f::<M>(i as f64)).collect(), ctx2);
1675        a += &b;
1676        let group0 = (1..=4).map(|i| f::<M>(i as f64));
1677        let group1 = (5..=8).map(|i| f::<M>(i as f64));
1678        let expected = group0
1679            .clone()
1680            .chain(group0)
1681            .chain(group1.clone())
1682            .chain(group1)
1683            .collect::<Vec<_>>();
1684        assert_eq!(triplet_values(&a), expected);
1685    }
1686
1687    /// The grouped flavour of [`test_batched_owned_lhs_add_broadcast_m`].
1688    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1689    pub fn test_grouped_owned_lhs_add_m<M>(ctx2: M::C)
1690    where
1691        M: Matrix + for<'a> Add<&'a M, Output = M>,
1692    {
1693        let wide_ctx = ctx4::<M>(&ctx2);
1694        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1695        let narrow = M::try_from_triplets(
1696            2,
1697            2,
1698            indices.clone(),
1699            (1..=8).map(|i| f::<M>(i as f64)).collect(),
1700            ctx2,
1701        )
1702        .unwrap();
1703        let wide = M::try_from_triplets(
1704            2,
1705            2,
1706            indices,
1707            (0..16).map(|_| f::<M>(0.0)).collect(),
1708            wide_ctx,
1709        )
1710        .unwrap();
1711        // the owned left-hand side governs, so the wide matrix is the owned one and the
1712        // narrow one broadcasts into its groups
1713        let c = wide + &narrow;
1714        assert_eq!(c.context().nbatch(), 4);
1715        let group0 = (1..=4).map(|i| f::<M>(i as f64));
1716        let group1 = (5..=8).map(|i| f::<M>(i as f64));
1717        let expected = group0
1718            .clone()
1719            .chain(group0)
1720            .chain(group1.clone())
1721            .chain(group1)
1722            .collect::<Vec<_>>();
1723        assert_eq!(triplet_values(&c), expected);
1724    }
1725
1726    /// The grouped flavour of [`test_batched_owned_lhs_narrow_m`].
1727    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1728    pub fn test_grouped_owned_lhs_narrow_m<M>(ctx2: M::C)
1729    where
1730        M: Matrix + for<'a> Add<&'a M, Output = M>,
1731    {
1732        let wide_ctx = ctx4::<M>(&ctx2);
1733        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
1734        let narrow = M::try_from_triplets(
1735            2,
1736            2,
1737            indices.clone(),
1738            (1..=8).map(|i| f::<M>(i as f64)).collect(),
1739            ctx2,
1740        )
1741        .unwrap();
1742        let wide = M::try_from_triplets(
1743            2,
1744            2,
1745            indices,
1746            (0..16).map(|_| f::<M>(0.0)).collect(),
1747            wide_ctx,
1748        )
1749        .unwrap();
1750        let _ = narrow + &wide;
1751    }
1752
1753    // --- Narrow-destination tests: a wider source would silently drop batches ---
1754
1755    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1756    pub fn test_narrow_dest_gemv_m<M: Matrix>(ctx2: M::C) {
1757        let wide = ctx4::<M>(&ctx2);
1758        let indices = vec![(0, 0), (1, 1)];
1759        let values = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1760        let a = M::try_from_triplets(2, 2, indices, values, wide.clone()).unwrap();
1761        let x = M::V::zeros(2, wide);
1762        let mut y = M::V::zeros(2, ctx2);
1763        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1764    }
1765
1766    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1767    pub fn test_narrow_dest_copy_from_m<M: Matrix>(ctx2: M::C) {
1768        let wide = ctx4::<M>(&ctx2);
1769        let indices = vec![(0, 0), (1, 1)];
1770        let values = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1771        let a = M::try_from_triplets(2, 2, indices, values, wide).unwrap();
1772        let mut b = M::zeros(2, 2, ctx2);
1773        b.copy_from(&a);
1774    }
1775
1776    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1777    pub fn test_narrow_dest_scale_add_m<M: Matrix>(ctx2: M::C) {
1778        let wide = ctx4::<M>(&ctx2);
1779        let indices = vec![(0, 0), (1, 1)];
1780        let vals = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1781        let x = M::try_from_triplets(2, 2, indices.clone(), vals.clone(), wide.clone()).unwrap();
1782        let y = M::try_from_triplets(2, 2, indices.clone(), vals, wide).unwrap();
1783        let dst = (1..=4).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1784        let mut result = M::try_from_triplets(2, 2, indices, dst, ctx2).unwrap();
1785        result.scale_add_and_assign(&x, f::<M>(2.0), &y);
1786    }
1787
1788    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1789    pub fn test_narrow_dest_gather_m<M: Matrix>(ctx2: M::C) {
1790        let wide = ctx4::<M>(&ctx2);
1791        let indices = vec![(0, 0), (1, 1)];
1792        let values = (1..=8).map(|i| f::<M>(i as f64)).collect::<Vec<_>>();
1793        let src = M::try_from_triplets(2, 2, indices.clone(), values, wide).unwrap();
1794        let zeros = (0..4).map(|_| f::<M>(0.0)).collect::<Vec<_>>();
1795        let mut dst = M::try_from_triplets(2, 2, indices, zeros, ctx2).unwrap();
1796        let gather_indices = <M::V as Vector>::Index::from_vec(vec![0, 1], Default::default());
1797        dst.gather(&src, &gather_indices);
1798    }
1799
1800    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1801    pub fn test_narrow_dest_set_data_with_indices_m<M: Matrix>(ctx2: M::C) {
1802        let wide = ctx4::<M>(&ctx2);
1803        let indices = vec![(0, 0), (1, 1)];
1804        let zeros = (0..4).map(|_| f::<M>(0.0)).collect::<Vec<_>>();
1805        let mut mat = M::try_from_triplets(2, 2, indices, zeros, ctx2).unwrap();
1806        let dst_indices = <M::V as Vector>::Index::from_vec(vec![0, 1], Default::default());
1807        let src_indices = <M::V as Vector>::Index::from_vec(vec![0, 1], Default::default());
1808        let data = M::V::from_vec((1..=8).map(|i| f::<M>(i as f64)).collect(), wide);
1809        mat.set_data_with_indices(&dst_indices, &src_indices, &data);
1810    }
1811
1812    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1813    pub fn test_narrow_dest_add_assign_m<M: DenseMatrix>(ctx2: M::C) {
1814        let wide = ctx4::<M>(&ctx2);
1815        let src = M::from_vec(2, 2, (0..16).map(|_| f::<M>(1.0)).collect(), wide);
1816        let mut dst = M::zeros(2, 2, ctx2);
1817        dst += &src;
1818    }
1819
1820    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1821    pub fn test_narrow_dest_gemv_cols_m<M: DenseMatrix>(ctx2: M::C) {
1822        let wide = ctx4::<M>(&ctx2);
1823        let a = M::from_vec(2, 2, (0..16).map(|_| f::<M>(1.0)).collect(), wide);
1824        let mut y = M::V::zeros(2, ctx2);
1825        let ones = vec![f::<M>(1.0), f::<M>(1.0)];
1826        a.gemv_cols(0, 2, f::<M>(1.0), &ones, f::<M>(0.0), &mut y);
1827    }
1828
1829    /// `gemv_cols` shares its coefficients across batches, so only the matrix broadcasts.
1830    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1831    pub fn test_grouped_gemv_cols_m<M: DenseMatrix>(ctx2: M::C) {
1832        let wide = ctx4::<M>(&ctx2);
1833        // batch 0 columns: [1, 1], [0, 0]; batch 1 columns: [2, 2], [0, 0]
1834        let a = M::from_vec(
1835            2,
1836            2,
1837            vec![
1838                f::<M>(1.0),
1839                f::<M>(1.0),
1840                f::<M>(0.0),
1841                f::<M>(0.0),
1842                f::<M>(2.0),
1843                f::<M>(2.0),
1844                f::<M>(0.0),
1845                f::<M>(0.0),
1846            ],
1847            ctx2,
1848        );
1849        let mut y = M::V::zeros(2, wide);
1850        let ones = vec![f::<M>(1.0), f::<M>(1.0)];
1851        a.gemv_cols(0, 2, f::<M>(1.0), &ones, f::<M>(0.0), &mut y);
1852        assert_eq!(
1853            y.clone_as_vec(),
1854            vec![
1855                f::<M>(1.0),
1856                f::<M>(1.0),
1857                f::<M>(1.0),
1858                f::<M>(1.0),
1859                f::<M>(2.0),
1860                f::<M>(2.0),
1861                f::<M>(2.0),
1862                f::<M>(2.0),
1863            ]
1864        );
1865    }
1866
1867    // --- Batched DenseMatrix-specific tests ---
1868
1869    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1870    pub fn test_batched_set_index_batch<M: DenseMatrix>(ctx: M::C) {
1871        assert_eq!(ctx.nbatch(), 2);
1872        // 2x2 matrix, nbatch=2: every (batch, row, col) gets a value that identifies it
1873        let mut a = M::zeros(2, 2, ctx);
1874        let val = |b: usize, i: usize, j: usize| f::<M>((100 * b + 10 * i + j) as f64);
1875        for b in 0..2 {
1876            for i in 0..2 {
1877                for j in 0..2 {
1878                    a.set_index_batch(b, i, j, val(b, i, j));
1879                }
1880            }
1881        }
1882        for b in 0..2 {
1883            for i in 0..2 {
1884                for j in 0..2 {
1885                    assert_eq!(mat_at(&a, b, i, j), val(b, i, j), "batch {b}, ({i}, {j})");
1886                }
1887            }
1888        }
1889    }
1890
1891    pub fn test_batched_from_vec<M: DenseMatrix>(ctx: M::C) {
1892        assert_eq!(ctx.nbatch(), 2);
1893        // 2x2 matrix, nbatch=2: physical is 2x4
1894        // batch0: col0=[1,3], col1=[2,4]; batch1: col0=[5,7], col1=[6,8]
1895        let a = M::from_vec(
1896            2,
1897            2,
1898            vec![
1899                f::<M>(1.0),
1900                f::<M>(3.0),
1901                f::<M>(2.0),
1902                f::<M>(4.0),
1903                f::<M>(5.0),
1904                f::<M>(7.0),
1905                f::<M>(6.0),
1906                f::<M>(8.0),
1907            ],
1908            ctx,
1909        );
1910        assert_eq!(a.nrows(), 2);
1911        assert_eq!(a.ncols(), 2);
1912        assert_eq!(at(&a, 0, 0, 0), f::<M>(1.0));
1913        assert_eq!(at(&a, 0, 1, 0), f::<M>(3.0));
1914        assert_eq!(at(&a, 0, 0, 1), f::<M>(2.0));
1915        assert_eq!(at(&a, 0, 1, 1), f::<M>(4.0));
1916        assert_eq!(at(&a, 1, 0, 0), f::<M>(5.0));
1917        assert_eq!(at(&a, 1, 1, 0), f::<M>(7.0));
1918        assert_eq!(at(&a, 1, 0, 1), f::<M>(6.0));
1919        assert_eq!(at(&a, 1, 1, 1), f::<M>(8.0));
1920    }
1921
1922    // --- Broadcasting tests ---
1923
1924    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1925    pub fn test_batched_gemv_o_broadcast_x<M: DenseMatrix>(ctx: M::C) {
1926        assert_eq!(ctx.nbatch(), 2);
1927        // 2x3 diff matrix, nbatch=2
1928        let diff = M::from_vec(
1929            2,
1930            3,
1931            vec![
1932                f::<M>(1.0),
1933                f::<M>(4.0),
1934                f::<M>(2.0),
1935                f::<M>(5.0),
1936                f::<M>(3.0),
1937                f::<M>(6.0),
1938                f::<M>(7.0),
1939                f::<M>(10.0),
1940                f::<M>(8.0),
1941                f::<M>(11.0),
1942                f::<M>(9.0),
1943                f::<M>(12.0),
1944            ],
1945            ctx.clone(),
1946        );
1947        // x with nbatch=1, length=2 (broadcast)
1948        let x = vec![f::<M>(1.0), f::<M>(1.0)];
1949        let mut y = M::V::zeros(2, ctx);
1950        diff.gemv_cols(0, 2, f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1951        // batch0: [[1,2],[4,5]] * [1,1] = [3, 9]
1952        // batch1: [[7,8],[10,11]] * [1,1] = [15, 21]
1953        assert_eq!(
1954            y.clone_as_vec(),
1955            vec![f::<M>(3.0), f::<M>(9.0), f::<M>(15.0), f::<M>(21.0)]
1956        );
1957    }
1958
1959    // --- Incompatible batch tests ---
1960
1961    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1962    pub fn test_batched_gemv_incompatible<M: DenseMatrix>(ctx2: M::C, ctx3: M::C) {
1963        assert_eq!(ctx2.nbatch(), 2);
1964        assert_eq!(ctx3.nbatch(), 3);
1965        let a = M::zeros(2, 2, ctx2.clone());
1966        let x = M::V::zeros(2, ctx3);
1967        let mut y = M::V::zeros(2, ctx2);
1968        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
1969    }
1970
1971    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
1972    pub fn test_batched_resize_cols<M: DenseMatrix>(ctx: M::C) {
1973        assert_eq!(ctx.nbatch(), 2);
1974        // 2x2, nbatch=2: batch0=[[1,2],[3,4]], batch1=[[5,6],[7,8]]
1975        let mut a = M::from_vec(
1976            2,
1977            2,
1978            vec![
1979                f::<M>(1.0),
1980                f::<M>(3.0),
1981                f::<M>(2.0),
1982                f::<M>(4.0),
1983                f::<M>(5.0),
1984                f::<M>(7.0),
1985                f::<M>(6.0),
1986                f::<M>(8.0),
1987            ],
1988            ctx.clone(),
1989        );
1990        // grow to 3 columns
1991        a.resize_cols(3);
1992        assert_eq!(a.ncols(), 3);
1993        assert_eq!(a.nrows(), 2);
1994        // existing data preserved per batch
1995        assert_eq!(at(&a, 0, 0, 0), f::<M>(1.0));
1996        assert_eq!(at(&a, 0, 1, 0), f::<M>(3.0));
1997        assert_eq!(at(&a, 0, 0, 1), f::<M>(2.0));
1998        assert_eq!(at(&a, 0, 1, 1), f::<M>(4.0));
1999        assert_eq!(at(&a, 1, 0, 0), f::<M>(5.0));
2000        assert_eq!(at(&a, 1, 1, 0), f::<M>(7.0));
2001        assert_eq!(at(&a, 1, 0, 1), f::<M>(6.0));
2002        assert_eq!(at(&a, 1, 1, 1), f::<M>(8.0));
2003        // new column is zero in every batch
2004        assert_eq!(at(&a, 0, 0, 2), f::<M>(0.0));
2005        assert_eq!(at(&a, 0, 1, 2), f::<M>(0.0));
2006        assert_eq!(at(&a, 1, 0, 2), f::<M>(0.0));
2007        assert_eq!(at(&a, 1, 1, 2), f::<M>(0.0));
2008        // verify via gemv that batch 1 data is intact
2009        let x = M::V::from_vec(
2010            vec![
2011                f::<M>(1.0),
2012                f::<M>(0.0),
2013                f::<M>(0.0),
2014                f::<M>(1.0),
2015                f::<M>(0.0),
2016                f::<M>(0.0),
2017            ],
2018            ctx.clone(),
2019        );
2020        let mut y = M::V::zeros(2, ctx.clone());
2021        a.gemv(f::<M>(1.0), &x, f::<M>(0.0), &mut y);
2022        // batch0: col0=[1,3], x=[1,0,0] → [1,3]
2023        // batch1: col0=[5,7], x=[1,0,0] → [5,7]
2024        assert_eq!(
2025            y.clone_as_vec(),
2026            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(5.0), f::<M>(7.0)]
2027        );
2028
2029        // shrink to 1 column
2030        a.resize_cols(1);
2031        assert_eq!(a.ncols(), 1);
2032        assert_eq!(at(&a, 0, 0, 0), f::<M>(1.0));
2033        assert_eq!(at(&a, 0, 1, 0), f::<M>(3.0));
2034        assert_eq!(at(&a, 1, 0, 0), f::<M>(5.0));
2035        assert_eq!(at(&a, 1, 1, 0), f::<M>(7.0));
2036        // verify batch1 col0 via gemv
2037        let x2 = M::V::from_vec(vec![f::<M>(1.0), f::<M>(1.0)], ctx.clone());
2038        let mut y2 = M::V::zeros(2, ctx);
2039        a.gemv(f::<M>(1.0), &x2, f::<M>(0.0), &mut y2);
2040        // batch0: [[1],[3]] * [1] = [1,3], batch1: [[5],[7]] * [1] = [5,7]
2041        assert_eq!(
2042            y2.clone_as_vec(),
2043            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(5.0), f::<M>(7.0)]
2044        );
2045    }
2046
2047    // --- New unbatched Matrix-generic tests ---
2048
2049    pub fn test_mul_scalar<M: Matrix>() {
2050        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2051        let values = vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)];
2052        let a = M::try_from_triplets(2, 2, indices, values, Default::default()).unwrap();
2053        let result = a * Scale(f::<M>(2.0));
2054        let (_, vals) = result.triplet_iter();
2055        let vals: Vec<_> = vals.collect();
2056        assert_eq!(
2057            vals,
2058            vec![f::<M>(2.0), f::<M>(6.0), f::<M>(4.0), f::<M>(8.0)]
2059        );
2060    }
2061
2062    pub fn test_add_column_to_vector<M: Matrix>() {
2063        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2064        let values = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0)];
2065        let mat = M::try_from_triplets(2, 2, indices, values, Default::default()).unwrap();
2066        let mut v = M::V::zeros(2, Default::default());
2067        mat.add_column_to_vector(1, &mut v);
2068        assert_eq!(v.clone_as_vec(), vec![f::<M>(3.0), f::<M>(4.0)]);
2069    }
2070
2071    // --- New unbatched DenseMatrix-specific tests ---
2072
2073    pub fn test_add<M: DenseMatrix>() {
2074        let a = M::from_vec(
2075            2,
2076            2,
2077            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
2078            Default::default(),
2079        );
2080        let b = M::from_vec(
2081            2,
2082            2,
2083            vec![f::<M>(5.0), f::<M>(7.0), f::<M>(6.0), f::<M>(8.0)],
2084            Default::default(),
2085        );
2086        let result = a + &b;
2087        assert_eq!(result.get_index(0, 0), f::<M>(6.0));
2088        assert_eq!(result.get_index(1, 1), f::<M>(12.0));
2089    }
2090
2091    pub fn test_sub<M: DenseMatrix>() {
2092        let a = M::from_vec(
2093            2,
2094            2,
2095            vec![f::<M>(5.0), f::<M>(7.0), f::<M>(6.0), f::<M>(8.0)],
2096            Default::default(),
2097        );
2098        let b = M::from_vec(
2099            2,
2100            2,
2101            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
2102            Default::default(),
2103        );
2104        let result = a - &b;
2105        assert_eq!(result.get_index(0, 0), f::<M>(4.0));
2106        assert_eq!(result.get_index(1, 1), f::<M>(4.0));
2107    }
2108
2109    pub fn test_add_assign<M: DenseMatrix>() {
2110        let mut a = M::from_vec(
2111            2,
2112            2,
2113            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
2114            Default::default(),
2115        );
2116        let b = M::from_vec(
2117            2,
2118            2,
2119            vec![f::<M>(5.0), f::<M>(7.0), f::<M>(6.0), f::<M>(8.0)],
2120            Default::default(),
2121        );
2122        a += &b;
2123        assert_eq!(a.get_index(0, 0), f::<M>(6.0));
2124        assert_eq!(a.get_index(1, 1), f::<M>(12.0));
2125    }
2126
2127    pub fn test_sub_assign<M: DenseMatrix>() {
2128        let mut a = M::from_vec(
2129            2,
2130            2,
2131            vec![f::<M>(5.0), f::<M>(7.0), f::<M>(6.0), f::<M>(8.0)],
2132            Default::default(),
2133        );
2134        let b = M::from_vec(
2135            2,
2136            2,
2137            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
2138            Default::default(),
2139        );
2140        a -= &b;
2141        assert_eq!(a.get_index(0, 0), f::<M>(4.0));
2142        assert_eq!(a.get_index(1, 1), f::<M>(4.0));
2143    }
2144
2145    /// `MatrixCommon::inner_mut` on an owned matrix; `test_matrix_common_by_ref` only reaches
2146    /// the generic `&M`/`&mut M` blanket impls.
2147    pub fn test_inner_mut<M: DenseMatrix>() {
2148        let mut a = M::zeros(2, 2, Default::default());
2149        let _ = a.inner_mut();
2150    }
2151
2152    /// `column_mut` on an unbatched (nbatch == 1) matrix: existing coverage only calls it through
2153    /// batched helpers, which take a different code path than the single-batch fast path.
2154    pub fn test_column_mut<M: DenseMatrix>() {
2155        let mut a = M::from_vec(
2156            2,
2157            2,
2158            vec![f::<M>(1.0), f::<M>(3.0), f::<M>(2.0), f::<M>(4.0)],
2159            Default::default(),
2160        );
2161        a.column_mut(1).set_index(0, f::<M>(20.0));
2162        assert_eq!(a.get_index(0, 1), f::<M>(20.0));
2163        assert_eq!(a.get_index(1, 1), f::<M>(4.0));
2164    }
2165
2166    pub fn test_gather<M: DenseMatrix>() {
2167        let mat1 = M::from_vec(
2168            3,
2169            3,
2170            vec![
2171                f::<M>(1.0),
2172                f::<M>(2.0),
2173                f::<M>(3.0),
2174                f::<M>(4.0),
2175                f::<M>(5.0),
2176                f::<M>(6.0),
2177                f::<M>(7.0),
2178                f::<M>(8.0),
2179                f::<M>(9.0),
2180            ],
2181            Default::default(),
2182        );
2183        let mut mat2 = M::zeros(2, 2, Default::default());
2184        let indices = <M::V as Vector>::Index::from_vec(vec![0, 1, 3, 4], Default::default());
2185        mat2.gather(&mat1, &indices);
2186        assert_eq!(mat2.get_index(0, 0), f::<M>(1.0));
2187        assert_eq!(mat2.get_index(1, 0), f::<M>(2.0));
2188        assert_eq!(mat2.get_index(0, 1), f::<M>(4.0));
2189        assert_eq!(mat2.get_index(1, 1), f::<M>(5.0));
2190    }
2191
2192    pub fn test_set_data_with_indices<M: DenseMatrix>() {
2193        let mut mat = M::zeros(2, 2, Default::default());
2194        let dst_indices = <M::V as Vector>::Index::from_vec(vec![0, 3], Default::default());
2195        let src_indices = <M::V as Vector>::Index::from_vec(vec![0, 1], Default::default());
2196        let data = M::V::from_vec(vec![f::<M>(5.0), f::<M>(6.0)], Default::default());
2197        mat.set_data_with_indices(&dst_indices, &src_indices, &data);
2198        assert_eq!(mat.get_index(0, 0), f::<M>(5.0));
2199        assert_eq!(mat.get_index(1, 1), f::<M>(6.0));
2200    }
2201
2202    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2203    pub fn test_batched_combine<M: DenseMatrix>(ctx: M::C) {
2204        assert_eq!(ctx.nbatch(), 2);
2205        #[rustfmt::skip]
2206        let data: Vec<M::T> = vec![
2207            // batch 0: 4x4 column-major (cols 0-3)
2208            f::<M>(1.0), f::<M>(2.0), f::<M>(3.0), f::<M>(4.0),
2209            f::<M>(5.0), f::<M>(6.0), f::<M>(7.0), f::<M>(8.0),
2210            f::<M>(9.0), f::<M>(10.0), f::<M>(11.0), f::<M>(12.0),
2211            f::<M>(13.0), f::<M>(14.0), f::<M>(15.0), f::<M>(16.0),
2212            // batch 1: 4x4 column-major (cols 0-3)
2213            f::<M>(101.0), f::<M>(102.0), f::<M>(103.0), f::<M>(104.0),
2214            f::<M>(105.0), f::<M>(106.0), f::<M>(107.0), f::<M>(108.0),
2215            f::<M>(109.0), f::<M>(110.0), f::<M>(111.0), f::<M>(112.0),
2216            f::<M>(113.0), f::<M>(114.0), f::<M>(115.0), f::<M>(116.0),
2217        ];
2218        let m = M::from_vec(4, 4, data, ctx.clone());
2219
2220        let alg_indices = <M::V as Vector>::Index::from_vec(vec![1, 3], Default::default());
2221        let [(ul, _), (ur, _), (ll, _), (lr, _)] = m.split(&alg_indices);
2222
2223        let recombined = M::combine(&ul, &ur, &ll, &lr, &alg_indices);
2224
2225        let (_orig_idx, orig_vals) = m.triplet_iter();
2226        let (_recom_idx, recom_vals) = recombined.triplet_iter();
2227        let orig_vals: Vec<_> = orig_vals.collect();
2228        let recom_vals: Vec<_> = recom_vals.collect();
2229        assert_eq!(orig_vals, recom_vals);
2230    }
2231
2232    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2233    pub fn test_batched_add_column_to_vector_m<M: Matrix>(ctx: M::C) {
2234        assert_eq!(ctx.nbatch(), 2);
2235        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2236        let values = vec![
2237            f::<M>(1.0),
2238            f::<M>(2.0),
2239            f::<M>(3.0),
2240            f::<M>(4.0),
2241            f::<M>(5.0),
2242            f::<M>(6.0),
2243            f::<M>(7.0),
2244            f::<M>(8.0),
2245        ];
2246        let mat = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
2247        let mut v = M::V::zeros(2, ctx);
2248        mat.add_column_to_vector(1, &mut v);
2249        assert_eq!(
2250            v.clone_as_vec(),
2251            vec![f::<M>(3.0), f::<M>(4.0), f::<M>(7.0), f::<M>(8.0)]
2252        );
2253    }
2254
2255    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2256    pub fn test_batched_add_columns_to_batched_vector_m<M: Matrix>(ctx: M::C) {
2257        assert_eq!(ctx.nbatch(), 2);
2258        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2259        let values = vec![
2260            f::<M>(1.0),
2261            f::<M>(2.0),
2262            f::<M>(3.0),
2263            f::<M>(4.0),
2264            f::<M>(5.0),
2265            f::<M>(6.0),
2266            f::<M>(7.0),
2267            f::<M>(8.0),
2268        ];
2269        let mat = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
2270        // one lane per (batch, column): lane `b * ncols + j` gets column `j` of batch `b`
2271        let ctx4 = ctx.clone_with_nbatch(4).unwrap();
2272        let mut v = M::V::zeros(2, ctx4);
2273        mat.add_columns_to_batched_vector(&mut v);
2274        assert_eq!(
2275            v.clone_as_vec(),
2276            vec![
2277                f::<M>(1.0),
2278                f::<M>(2.0),
2279                f::<M>(3.0),
2280                f::<M>(4.0),
2281                f::<M>(5.0),
2282                f::<M>(6.0),
2283                f::<M>(7.0),
2284                f::<M>(8.0)
2285            ]
2286        );
2287        // it accumulates, like its single-column sibling
2288        mat.add_columns_to_batched_vector(&mut v);
2289        assert_eq!(v.clone_as_vec()[0], f::<M>(2.0));
2290        assert_eq!(v.clone_as_vec()[7], f::<M>(16.0));
2291    }
2292
2293    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2294    pub fn test_batched_set_data_with_indices_m<M: Matrix>(ctx: M::C) {
2295        assert_eq!(ctx.nbatch(), 2);
2296        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2297        let zero_values = vec![
2298            f::<M>(0.0),
2299            f::<M>(0.0),
2300            f::<M>(0.0),
2301            f::<M>(0.0),
2302            f::<M>(0.0),
2303            f::<M>(0.0),
2304            f::<M>(0.0),
2305            f::<M>(0.0),
2306        ];
2307        let mut mat = M::try_from_triplets(2, 2, indices, zero_values, ctx.clone()).unwrap();
2308        let dst_indices = <M::V as Vector>::Index::from_vec(vec![0, 3], Default::default());
2309        let src_indices = <M::V as Vector>::Index::from_vec(vec![0, 1], Default::default());
2310        let data = M::V::from_vec(
2311            vec![f::<M>(5.0), f::<M>(6.0), f::<M>(50.0), f::<M>(60.0)],
2312            ctx,
2313        );
2314        mat.set_data_with_indices(&dst_indices, &src_indices, &data);
2315        let (_, vals) = mat.triplet_iter();
2316        let vals: Vec<_> = vals.collect();
2317        assert_eq!(
2318            vals,
2319            vec![
2320                f::<M>(5.0),
2321                f::<M>(0.0),
2322                f::<M>(0.0),
2323                f::<M>(6.0),
2324                f::<M>(50.0),
2325                f::<M>(0.0),
2326                f::<M>(0.0),
2327                f::<M>(60.0),
2328            ]
2329        );
2330    }
2331
2332    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2333    pub fn test_batched_gather_m<M: Matrix>(ctx: M::C) {
2334        assert_eq!(ctx.nbatch(), 2);
2335        let indices: Vec<(IndexType, IndexType)> =
2336            (0..3).flat_map(|j| (0..3).map(move |i| (i, j))).collect();
2337        let values = vec![
2338            f::<M>(1.0),
2339            f::<M>(2.0),
2340            f::<M>(3.0),
2341            f::<M>(4.0),
2342            f::<M>(5.0),
2343            f::<M>(6.0),
2344            f::<M>(7.0),
2345            f::<M>(8.0),
2346            f::<M>(9.0),
2347            f::<M>(10.0),
2348            f::<M>(20.0),
2349            f::<M>(30.0),
2350            f::<M>(40.0),
2351            f::<M>(50.0),
2352            f::<M>(60.0),
2353            f::<M>(70.0),
2354            f::<M>(80.0),
2355            f::<M>(90.0),
2356        ];
2357        let mat1 = M::try_from_triplets(3, 3, indices, values, ctx.clone()).unwrap();
2358        let dest_indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2359        let zero_values = vec![
2360            f::<M>(0.0),
2361            f::<M>(0.0),
2362            f::<M>(0.0),
2363            f::<M>(0.0),
2364            f::<M>(0.0),
2365            f::<M>(0.0),
2366            f::<M>(0.0),
2367            f::<M>(0.0),
2368        ];
2369        let mut mat2 = M::try_from_triplets(2, 2, dest_indices, zero_values, ctx).unwrap();
2370        let gather_indices =
2371            <M::V as Vector>::Index::from_vec(vec![0, 1, 3, 4], Default::default());
2372        mat2.gather(&mat1, &gather_indices);
2373        let (_, vals) = mat2.triplet_iter();
2374        let vals: Vec<_> = vals.collect();
2375        assert_eq!(
2376            vals,
2377            vec![
2378                f::<M>(1.0),
2379                f::<M>(2.0),
2380                f::<M>(4.0),
2381                f::<M>(5.0),
2382                f::<M>(10.0),
2383                f::<M>(20.0),
2384                f::<M>(40.0),
2385                f::<M>(50.0),
2386            ]
2387        );
2388    }
2389
2390    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2391    pub fn test_batched_mul_scalar_m<M: Matrix>(ctx: M::C) {
2392        assert_eq!(ctx.nbatch(), 2);
2393        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2394        let values = vec![
2395            f::<M>(1.0),
2396            f::<M>(3.0),
2397            f::<M>(2.0),
2398            f::<M>(4.0),
2399            f::<M>(5.0),
2400            f::<M>(7.0),
2401            f::<M>(6.0),
2402            f::<M>(8.0),
2403        ];
2404        let a = M::try_from_triplets(2, 2, indices, values, ctx.clone()).unwrap();
2405        let result = a * Scale(f::<M>(2.0));
2406        let (_, vals) = result.triplet_iter();
2407        let vals: Vec<_> = vals.collect();
2408        assert_eq!(
2409            vals,
2410            vec![
2411                f::<M>(2.0),
2412                f::<M>(6.0),
2413                f::<M>(4.0),
2414                f::<M>(8.0),
2415                f::<M>(10.0),
2416                f::<M>(14.0),
2417                f::<M>(12.0),
2418                f::<M>(16.0),
2419            ]
2420        );
2421    }
2422
2423    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2424    pub fn test_batched_partition_indices<M: Matrix>(ctx: M::C) {
2425        assert_eq!(ctx.nbatch(), 2);
2426        let zero_val = M::T::zero();
2427        let one_val = f::<M>(1.0);
2428        let two_val = f::<M>(2.0);
2429        let indices = vec![(0, 0), (1, 1), (2, 2)];
2430        let values = vec![one_val, zero_val, one_val, two_val, zero_val, two_val];
2431        let m = M::try_from_triplets(3, 3, indices, values, ctx).unwrap();
2432        let (zero_idx, nonzero_idx) = m.partition_indices_by_zero_diagonal();
2433        assert_eq!(zero_idx.clone_as_vec(), vec![1]);
2434        assert_eq!(nonzero_idx.clone_as_vec(), vec![0, 2]);
2435    }
2436
2437    /// Same order-1 scan as [`test_update_backward_diff`], run on 2 independent batches with
2438    /// distinct data, verified via `column(..).clone_as_vec()` so both batches are checked
2439    /// (`get_index` only ever reads batch 0).
2440    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2441    pub fn test_batched_update_backward_diff<M: DenseMatrix>(ctx: M::C) {
2442        assert_eq!(ctx.nbatch(), 2);
2443        let order = 1;
2444        let mut a = M::from_vec(
2445            2,
2446            4,
2447            vec![
2448                // batch 0
2449                f::<M>(1.0),
2450                f::<M>(10.0),
2451                f::<M>(2.0),
2452                f::<M>(20.0),
2453                f::<M>(3.0),
2454                f::<M>(30.0),
2455                f::<M>(0.0),
2456                f::<M>(0.0),
2457                // batch 1
2458                f::<M>(5.0),
2459                f::<M>(50.0),
2460                f::<M>(6.0),
2461                f::<M>(60.0),
2462                f::<M>(7.0),
2463                f::<M>(70.0),
2464                f::<M>(0.0),
2465                f::<M>(0.0),
2466            ],
2467            ctx.clone(),
2468        );
2469        let d = M::V::from_vec(
2470            vec![f::<M>(100.0), f::<M>(1000.0), f::<M>(200.0), f::<M>(2000.0)],
2471            ctx,
2472        );
2473
2474        a.update_backward_diff(order, &d);
2475
2476        assert_eq!(
2477            a.column(3).into_owned().clone_as_vec(),
2478            vec![f::<M>(97.0), f::<M>(970.0), f::<M>(193.0), f::<M>(1930.0)]
2479        );
2480        assert_eq!(
2481            a.column(2).into_owned().clone_as_vec(),
2482            vec![f::<M>(100.0), f::<M>(1000.0), f::<M>(200.0), f::<M>(2000.0)]
2483        );
2484        assert_eq!(
2485            a.column(1).into_owned().clone_as_vec(),
2486            vec![f::<M>(102.0), f::<M>(1020.0), f::<M>(206.0), f::<M>(2060.0)]
2487        );
2488        assert_eq!(
2489            a.column(0).into_owned().clone_as_vec(),
2490            vec![f::<M>(103.0), f::<M>(1030.0), f::<M>(211.0), f::<M>(2110.0)]
2491        );
2492    }
2493
2494    #[cfg_attr(not(feature = "cuda"), allow(dead_code))]
2495    pub fn test_batched_from_diagonal_dense<M: DenseMatrix>(ctx: M::C) {
2496        assert_eq!(ctx.nbatch(), 2);
2497        let v = M::V::from_vec(
2498            vec![f::<M>(2.0), f::<M>(3.0), f::<M>(4.0), f::<M>(5.0)],
2499            ctx,
2500        );
2501        let a = M::from_diagonal(&v);
2502        assert_eq!(a.nrows(), 2);
2503        assert_eq!(a.ncols(), 2);
2504        assert_eq!(at(&a, 0, 0, 0), f::<M>(2.0));
2505        assert_eq!(at(&a, 0, 1, 1), f::<M>(3.0));
2506        assert_eq!(at(&a, 0, 0, 1), f::<M>(0.0));
2507        assert_eq!(at(&a, 0, 1, 0), f::<M>(0.0));
2508        assert_eq!(at(&a, 1, 0, 0), f::<M>(4.0));
2509        assert_eq!(at(&a, 1, 1, 1), f::<M>(5.0));
2510        assert_eq!(at(&a, 1, 0, 1), f::<M>(0.0));
2511        assert_eq!(at(&a, 1, 1, 0), f::<M>(0.0));
2512    }
2513
2514    pub fn test_try_from_triplets_wrong_length<M: Matrix>() {
2515        let indices = vec![(0, 0), (1, 0), (0, 1), (1, 1)];
2516        // one value too few: triggers the length assertion inside try_from_triplets
2517        let values = vec![f::<M>(1.0), f::<M>(2.0), f::<M>(3.0)];
2518        let _ = M::try_from_triplets(2, 2, indices, values, Default::default());
2519    }
2520
2521    // --- Batched view-mut tests ---
2522}
2523
2524#[cfg(test)]
2525macro_rules! generate_matrix_tests_nonbatched {
2526    ($suffix:ident, $M:ty) => {
2527        paste::paste! {
2528            #[test]
2529            fn [<test_zeros_ $suffix>]() {
2530                $crate::matrix::tests::test_zeros::<$M>();
2531            }
2532            #[test]
2533            fn [<test_sparsity_ $suffix>]() {
2534                $crate::matrix::tests::test_sparsity::<$M>();
2535            }
2536            #[test]
2537            fn [<test_matrix_common_by_ref_ $suffix>]() {
2538                $crate::matrix::tests::test_matrix_common_by_ref::<$M>();
2539            }
2540            #[test]
2541            fn [<test_from_diagonal_ $suffix>]() {
2542                $crate::matrix::tests::test_from_diagonal::<$M>();
2543            }
2544            #[test]
2545            fn [<test_gemv_ $suffix>]() {
2546                $crate::matrix::tests::test_gemv::<$M>();
2547            }
2548            #[test]
2549            fn [<test_set_column_ $suffix>]() {
2550                $crate::matrix::tests::test_set_column::<$M>();
2551            }
2552            #[test]
2553            fn [<test_copy_from_ $suffix>]() {
2554                $crate::matrix::tests::test_copy_from::<$M>();
2555            }
2556            #[test]
2557            fn [<test_scale_add_and_assign_ $suffix>]() {
2558                $crate::matrix::tests::test_scale_add_and_assign::<$M>();
2559            }
2560            #[test]
2561            fn [<test_partition_indices_ $suffix>]() {
2562                $crate::matrix::tests::test_partition_indices_by_zero_diagonal::<$M>();
2563            }
2564            #[test]
2565            fn [<test_mul_scalar_ $suffix>]() {
2566                $crate::matrix::tests::test_mul_scalar::<$M>();
2567            }
2568            #[test]
2569            fn [<test_add_column_to_vector_ $suffix>]() {
2570                $crate::matrix::tests::test_add_column_to_vector::<$M>();
2571            }
2572            #[test]
2573            #[should_panic]
2574            fn [<test_try_from_triplets_wrong_length_ $suffix>]() {
2575                $crate::matrix::tests::test_try_from_triplets_wrong_length::<$M>();
2576            }
2577        }
2578    };
2579}
2580
2581#[cfg(test)]
2582#[cfg_attr(not(feature = "cuda"), allow(unused_macros))]
2583macro_rules! generate_matrix_tests_batched {
2584    ($suffix:ident, $M:ty, $ctx1:expr, $ctx2:expr) => {
2585        paste::paste! {
2586            #[test]
2587            fn [<test_batched_add_column_to_vector_ $suffix>]() {
2588                $crate::matrix::tests::test_batched_add_column_to_vector_m::<$M>($ctx2);
2589            }
2590            #[test]
2591            fn [<test_batched_add_columns_to_batched_vector_ $suffix>]() {
2592                $crate::matrix::tests::test_batched_add_columns_to_batched_vector_m::<$M>($ctx2);
2593            }
2594            #[test]
2595            fn [<test_batched_set_data_with_indices_ $suffix>]() {
2596                $crate::matrix::tests::test_batched_set_data_with_indices_m::<$M>($ctx2);
2597            }
2598            #[test]
2599            fn [<test_batched_gather_ $suffix>]() {
2600                $crate::matrix::tests::test_batched_gather_m::<$M>($ctx2);
2601            }
2602            #[test]
2603            fn [<test_batched_mul_scalar_ $suffix>]() {
2604                $crate::matrix::tests::test_batched_mul_scalar_m::<$M>($ctx2);
2605            }
2606            #[test]
2607            fn [<test_batched_partition_indices_ $suffix>]() {
2608                $crate::matrix::tests::test_batched_partition_indices::<$M>($ctx2);
2609            }
2610            #[test]
2611            fn [<test_batched_zeros_ $suffix>]() {
2612                $crate::matrix::tests::test_batched_zeros_m::<$M>($ctx2);
2613            }
2614            #[test]
2615            fn [<test_batched_gemv_ $suffix>]() {
2616                $crate::matrix::tests::test_batched_gemv_m::<$M>($ctx2);
2617            }
2618            #[test]
2619            fn [<test_batched_gemv_broadcast_x_ $suffix>]() {
2620                $crate::matrix::tests::test_batched_gemv_broadcast_x_m::<$M>($ctx2);
2621            }
2622            #[test]
2623            fn [<test_batched_gemv_broadcast_mat_ $suffix>]() {
2624                $crate::matrix::tests::test_batched_gemv_broadcast_mat_m::<$M>($ctx2);
2625            }
2626            #[test]
2627            fn [<test_batched_from_diagonal_ $suffix>]() {
2628                $crate::matrix::tests::test_batched_from_diagonal_m::<$M>($ctx2);
2629            }
2630            #[test]
2631            fn [<test_batched_copy_from_ $suffix>]() {
2632                $crate::matrix::tests::test_batched_copy_from_m::<$M>($ctx2);
2633            }
2634            #[test]
2635            fn [<test_batched_copy_from_broadcast_ $suffix>]() {
2636                $crate::matrix::tests::test_batched_copy_from_broadcast_m::<$M>($ctx2);
2637            }
2638            #[test]
2639            fn [<test_batched_set_column_ $suffix>]() {
2640                $crate::matrix::tests::test_batched_set_column_m::<$M>($ctx2);
2641            }
2642            #[test]
2643            fn [<test_batched_scale_add_ $suffix>]() {
2644                $crate::matrix::tests::test_batched_scale_add_and_assign_m::<$M>($ctx2);
2645            }
2646            #[test]
2647            fn [<test_batched_owned_lhs_broadcast_ $suffix>]() {
2648                $crate::matrix::tests::test_batched_owned_lhs_broadcast_m::<$M>($ctx2);
2649            }
2650            #[test]
2651            fn [<test_batched_scale_add_broadcast_ $suffix>]() {
2652                $crate::matrix::tests::test_batched_scale_add_and_assign_broadcast_m::<$M>($ctx2);
2653            }
2654
2655            #[test]
2656            fn [<test_batched_mul_scalar_ref_ $suffix>]() {
2657                $crate::matrix::tests::test_batched_mul_scalar_ref_m::<$M>($ctx2);
2658            }
2659            #[test]
2660            fn [<test_batched_owned_lhs_add_broadcast_ $suffix>]() {
2661                $crate::matrix::tests::test_batched_owned_lhs_add_broadcast_m::<$M>($ctx2);
2662            }
2663            #[test]
2664            fn [<test_batched_gemv_beta_ $suffix>]() {
2665                $crate::matrix::tests::test_batched_gemv_beta_m::<$M>($ctx2);
2666            }
2667
2668            // --- Governing-operand tests: the operand that governs must be the wider one ---
2669            #[test]
2670            #[should_panic(expected = "incompatible nbatch")]
2671            fn [<test_batched_owned_lhs_narrow_ $suffix>]() {
2672                $crate::matrix::tests::test_batched_owned_lhs_narrow_m::<$M>($ctx2);
2673            }
2674            #[test]
2675            #[should_panic(expected = "incompatible nbatch")]
2676            fn [<test_grouped_owned_lhs_narrow_ $suffix>]() {
2677                $crate::matrix::tests::test_grouped_owned_lhs_narrow_m::<$M>($ctx2);
2678            }
2679
2680            // --- Narrow-destination tests (source wider than destination panics) ---
2681            #[test]
2682            #[should_panic(expected = "incompatible nbatch")]
2683            fn [<test_narrow_dest_gemv_ $suffix>]() {
2684                $crate::matrix::tests::test_narrow_dest_gemv_m::<$M>($ctx2);
2685            }
2686            #[test]
2687            #[should_panic(expected = "incompatible nbatch")]
2688            fn [<test_narrow_dest_copy_from_ $suffix>]() {
2689                $crate::matrix::tests::test_narrow_dest_copy_from_m::<$M>($ctx2);
2690            }
2691            #[test]
2692            #[should_panic(expected = "incompatible nbatch")]
2693            fn [<test_narrow_dest_scale_add_ $suffix>]() {
2694                $crate::matrix::tests::test_narrow_dest_scale_add_m::<$M>($ctx2);
2695            }
2696            #[test]
2697            #[should_panic(expected = "incompatible nbatch")]
2698            fn [<test_narrow_dest_gather_ $suffix>]() {
2699                $crate::matrix::tests::test_narrow_dest_gather_m::<$M>($ctx2);
2700            }
2701            #[test]
2702            #[should_panic(expected = "incompatible nbatch")]
2703            fn [<test_narrow_dest_set_data_with_indices_ $suffix>]() {
2704                $crate::matrix::tests::test_narrow_dest_set_data_with_indices_m::<$M>($ctx2);
2705            }
2706
2707            // --- Grouped broadcast tests (B -> B * P, using $ctx2 widened to 4) ---
2708            #[test]
2709            fn [<test_grouped_gemv_ $suffix>]() {
2710                $crate::matrix::tests::test_grouped_gemv_m::<$M>($ctx2);
2711            }
2712            #[test]
2713            fn [<test_grouped_gemv_x_ $suffix>]() {
2714                $crate::matrix::tests::test_grouped_gemv_x_m::<$M>($ctx2);
2715            }
2716            #[test]
2717            fn [<test_grouped_copy_from_ $suffix>]() {
2718                $crate::matrix::tests::test_grouped_copy_from_m::<$M>($ctx2);
2719            }
2720            #[test]
2721            fn [<test_grouped_owned_lhs_add_ $suffix>]() {
2722                $crate::matrix::tests::test_grouped_owned_lhs_add_m::<$M>($ctx2);
2723            }
2724            #[test]
2725            fn [<test_grouped_scale_add_ $suffix>]() {
2726                $crate::matrix::tests::test_grouped_scale_add_and_assign_m::<$M>($ctx2);
2727            }
2728        }
2729    };
2730}
2731
2732#[cfg(test)]
2733macro_rules! generate_dense_matrix_tests_nonbatched {
2734    ($suffix:ident, $M:ty) => {
2735        paste::paste! {
2736            #[test]
2737            fn [<test_from_vec_ $suffix>]() {
2738                $crate::matrix::tests::test_from_vec::<$M>();
2739            }
2740            #[test]
2741            fn [<test_from_diagonal_dense_ $suffix>]() {
2742                $crate::matrix::tests::test_from_diagonal_dense::<$M>();
2743            }
2744            #[test]
2745            fn [<test_column_view_ $suffix>]() {
2746                $crate::matrix::tests::test_column_view::<$M>();
2747            }
2748            #[test]
2749            fn [<test_sparsity_dense_errors_ $suffix>]() {
2750                $crate::matrix::tests::test_sparsity_dense_errors::<$M>();
2751            }
2752            #[test]
2753            fn [<test_update_backward_diff_ $suffix>]() {
2754                $crate::matrix::tests::test_update_backward_diff::<$M>();
2755            }
2756            #[test]
2757            fn [<test_gemv_cols_ $suffix>]() {
2758                $crate::matrix::tests::test_gemv_cols::<$M>();
2759            }
2760
2761            #[test]
2762            fn [<test_gemv_cols_many_columns_ $suffix>]() {
2763                $crate::matrix::tests::test_gemv_cols_many_columns::<$M>();
2764            }
2765
2766            #[test]
2767            fn [<test_mul_cols_by_ $suffix>]() {
2768                $crate::matrix::tests::test_mul_cols_by::<$M>();
2769            }
2770            #[test]
2771            fn [<test_resize_cols_ $suffix>]() {
2772                $crate::matrix::tests::test_resize_cols::<$M>();
2773            }
2774            #[test]
2775            fn [<test_add_ $suffix>]() {
2776                $crate::matrix::tests::test_add::<$M>();
2777            }
2778            #[test]
2779            fn [<test_sub_ $suffix>]() {
2780                $crate::matrix::tests::test_sub::<$M>();
2781            }
2782            #[test]
2783            fn [<test_inner_mut_ $suffix>]() {
2784                $crate::matrix::tests::test_inner_mut::<$M>();
2785            }
2786            #[test]
2787            fn [<test_column_mut_ $suffix>]() {
2788                $crate::matrix::tests::test_column_mut::<$M>();
2789            }
2790            #[test]
2791            fn [<test_add_assign_ $suffix>]() {
2792                $crate::matrix::tests::test_add_assign::<$M>();
2793            }
2794            #[test]
2795            fn [<test_sub_assign_ $suffix>]() {
2796                $crate::matrix::tests::test_sub_assign::<$M>();
2797            }
2798            #[test]
2799            fn [<test_gather_ $suffix>]() {
2800                $crate::matrix::tests::test_gather::<$M>();
2801            }
2802            #[test]
2803            fn [<test_set_data_with_indices_ $suffix>]() {
2804                $crate::matrix::tests::test_set_data_with_indices::<$M>();
2805            }
2806        }
2807    };
2808}
2809
2810#[cfg(test)]
2811#[cfg_attr(not(feature = "cuda"), allow(unused_macros))]
2812macro_rules! generate_dense_matrix_tests_batched {
2813    ($suffix:ident, $M:ty, $ctx1:expr, $ctx2:expr) => {
2814        paste::paste! {
2815            #[test]
2816            fn [<test_batched_update_backward_diff_ $suffix>]() {
2817                $crate::matrix::tests::test_batched_update_backward_diff::<$M>($ctx2);
2818            }
2819            #[test]
2820            fn [<test_batched_gemv_cols_ $suffix>]() {
2821                $crate::matrix::tests::test_batched_gemv_cols_m::<$M>($ctx2);
2822            }
2823
2824            #[test]
2825            fn [<test_batched_mul_cols_by_ $suffix>]() {
2826                $crate::matrix::tests::test_batched_mul_cols_by::<$M>($ctx2);
2827            }
2828            #[test]
2829            fn [<test_batched_from_diagonal_dense_ $suffix>]() {
2830                $crate::matrix::tests::test_batched_from_diagonal_dense::<$M>($ctx2);
2831            }
2832            #[test]
2833            fn [<test_batched_from_vec_ $suffix>]() {
2834                $crate::matrix::tests::test_batched_from_vec::<$M>($ctx2);
2835            }
2836            #[test]
2837            fn [<test_batched_set_index_batch_ $suffix>]() {
2838                $crate::matrix::tests::test_batched_set_index_batch::<$M>($ctx2);
2839            }
2840            #[test]
2841            fn [<test_batched_gemv_o_broadcast_x_ $suffix>]() {
2842                $crate::matrix::tests::test_batched_gemv_o_broadcast_x::<$M>($ctx2);
2843            }
2844            #[test]
2845            fn [<test_batched_resize_cols_ $suffix>]() {
2846                $crate::matrix::tests::test_batched_resize_cols::<$M>($ctx2);
2847            }
2848            #[test]
2849            fn [<test_batched_combine_ $suffix>]() {
2850                $crate::matrix::tests::test_batched_combine::<$M>($ctx2);
2851            }
2852            #[test]
2853            fn [<test_batched_add_assign_broadcast_ $suffix>]() {
2854                $crate::matrix::tests::test_batched_add_assign_broadcast_m::<$M>($ctx2);
2855            }
2856            #[test]
2857            fn [<test_grouped_add_assign_ $suffix>]() {
2858                $crate::matrix::tests::test_grouped_add_assign_m::<$M>($ctx2);
2859            }
2860            #[test]
2861            #[should_panic(expected = "incompatible nbatch")]
2862            fn [<test_narrow_dest_gemv_cols_ $suffix>]() {
2863                $crate::matrix::tests::test_narrow_dest_gemv_cols_m::<$M>($ctx2);
2864            }
2865            #[test]
2866            #[should_panic(expected = "incompatible nbatch")]
2867            fn [<test_narrow_dest_add_assign_ $suffix>]() {
2868                $crate::matrix::tests::test_narrow_dest_add_assign_m::<$M>($ctx2);
2869            }
2870            #[test]
2871            fn [<test_grouped_gemv_cols_ $suffix>]() {
2872                $crate::matrix::tests::test_grouped_gemv_cols_m::<$M>($ctx2);
2873            }
2874            #[test]
2875            #[should_panic(expected = "incompatible nbatch")]
2876            fn [<test_batched_gemv_incompatible_ $suffix>]() {
2877                $crate::matrix::tests::test_batched_gemv_incompatible::<$M>($ctx2, $ctx1.clone_with_nbatch(3).unwrap());
2878            }
2879        }
2880    };
2881}
2882
2883#[cfg(test)]
2884#[cfg_attr(not(feature = "cuda"), allow(unused_imports))]
2885pub(crate) use generate_dense_matrix_tests_batched;
2886#[cfg(test)]
2887pub(crate) use generate_dense_matrix_tests_nonbatched;
2888#[cfg(test)]
2889#[cfg_attr(not(feature = "cuda"), allow(unused_imports))]
2890pub(crate) use generate_matrix_tests_batched;
2891#[cfg(test)]
2892pub(crate) use generate_matrix_tests_nonbatched;