## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/trueno/src/matrix/ops/ml_ops.rs

    Line
    Count
    Source
    1
    //! Machine learning operations for Matrix
    2
    //!
    3
    //! This module provides ML-specific operations:
    4
    //! - `convolve2d()` - 2D convolution
    5
    //! - `embedding_lookup()` - Embedding table lookup
    6
    //! - `embedding_lookup_sparse()` - Embedding lookup with gradient tracking
    7
    //! - `max_pool2d()` - Max pooling
    8
    //! - `avg_pool2d()` - Average pooling
    9
    //! - `topk()` - Top-K selection
    10
    //! - `gather()` - Gather elements along axis
    11
    //! - `pad()` - Pad matrix with constant value
    12
    13
    use crate::TruenoError;
    14
    15
    use super::super::Matrix;
    16
    17
    impl Matrix<f32> {
    18
        /// Perform 2D convolution with a kernel
    19
        ///
    20
        /// Applies a 2D convolution operation using "valid" padding (no padding),
    21
        /// resulting in an output smaller than the input.
    22
        ///
    23
        /// # Arguments
    24
        ///
    25
        /// * `kernel` - Convolution kernel (filter) to apply
    26
        ///
    27
        /// # Returns
    28
        ///
    29
        /// Convolved matrix with dimensions:
    30
        /// - rows: `input.rows - kernel.rows + 1`
    31
        /// - cols: `input.cols - kernel.cols + 1`
    32
        ///
    33
        /// # Errors
    34
        ///
    35
        /// Returns `InvalidInput` if:
    36
        /// - Kernel is larger than input in any dimension
    37
        ///
    38
        /// # Example
    39
        ///
    40
        /// ```
    41
        /// use trueno::Matrix;
    42
        ///
    43
        /// // 5x5 input image
    44
        /// let input = Matrix::from_vec(
    45
        ///     5, 5,
    46
        ///     vec![
    47
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    48
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    49
        ///         0.0, 0.0, 9.0, 0.0, 0.0,
    50
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    51
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    52
        ///     ]
    53
        /// ).unwrap();
    54
        ///
    55
        /// // 3x3 averaging kernel
    56
        /// let kernel_val = 1.0 / 9.0;
    57
        /// let kernel = Matrix::from_vec(
    58
        ///     3, 3,
    59
        ///     vec![kernel_val; 9]
    60
        /// ).unwrap();
    61
        ///
    62
        /// let result = input.convolve2d(&kernel).unwrap();
    63
        /// assert_eq!(result.rows(), 3); // 5 - 3 + 1
    64
        /// assert_eq!(result.cols(), 3);
    65
        /// ```
    66
        // =========================================================================
    67
        // HOT PATH - PERFORMANCE CRITICAL
    68
        // =========================================================================
    69
        // This function processes millions of elements for typical image sizes.
    70
        // Any changes to the inner loop REQUIRE benchmark verification.
    71
        // =========================================================================
    72
    0
        pub fn convolve2d(&self, kernel: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    73
            // Validate kernel size
    74
    0
            if kernel.rows > self.rows || kernel.cols > self.cols {
    75
    0
                return Err(TruenoError::InvalidInput(format!(
    76
    0
                    "Kernel size ({}x{}) larger than input ({}x{})",
    77
    0
                    kernel.rows, kernel.cols, self.rows, self.cols
    78
    0
                )));
    79
    0
            }
    80
    81
            // Calculate output dimensions (valid padding)
    82
    0
            let output_rows = self.rows - kernel.rows + 1;
    83
    0
            let output_cols = self.cols - kernel.cols + 1;
    84
    85
            // Initialize output matrix (reuse parent's backend)
    86
    0
            let mut result = Matrix::zeros_with_backend(output_rows, output_cols, self.backend);
    87
    88
            // GPU acceleration for large convolutions
    89
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    90
            const GPU_THRESHOLD: usize = 10_000;
    91
    92
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    93
            {
    94
    0
                if output_rows * output_cols >= GPU_THRESHOLD {
    95
                    use crate::backends::gpu::GpuBackend;
    96
    97
    0
                    if GpuBackend::is_available() {
    98
    0
                        if let Ok(gpu_result) =
    99
    0
                            self.convolve2d_gpu(kernel, &mut result, output_rows, output_cols)
    100
                        {
    101
    0
                            return Ok(gpu_result);
    102
    0
                        }
    103
    0
                    }
    104
    0
                }
    105
            }
    106
    107
            // Scalar baseline implementation - optimized with direct indexing
    108
    0
            let input_data = self.as_slice();
    109
    0
            let kernel_data = kernel.as_slice();
    110
    0
            let result_data = result.data.as_mut_slice();
    111
    0
            let input_cols = self.cols;
    112
    0
            let kernel_cols = kernel.cols;
    113
    0
            let result_cols = output_cols;
    114
    115
    0
            for out_row in 0..output_rows {
    116
    0
                for out_col in 0..output_cols {
    117
    0
                    let mut sum = 0.0;
    118
    119
    0
                    for k_row in 0..kernel.rows {
    120
    0
                        let in_row = out_row + k_row;
    121
    0
                        let input_row_offset = in_row * input_cols;
    122
    0
                        let kernel_row_offset = k_row * kernel_cols;
    123
    124
    0
                        for k_col in 0..kernel.cols {
    125
    0
                            let in_col = out_col + k_col;
    126
    0
                            sum += input_data[input_row_offset + in_col]
    127
    0
                                * kernel_data[kernel_row_offset + k_col];
    128
    0
                        }
    129
                    }
    130
    131
    0
                    result_data[out_row * result_cols + out_col] = sum;
    132
                }
    133
            }
    134
    135
    0
            Ok(result)
    136
    0
        }
    137
    138
        /// GPU-accelerated 2D convolution helper
    139
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    140
    0
        fn convolve2d_gpu(
    141
    0
            &self,
    142
    0
            kernel: &Matrix<f32>,
    143
    0
            result: &mut Matrix<f32>,
    144
    0
            _output_rows: usize,
    145
    0
            _output_cols: usize,
    146
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    147
            use crate::backends::gpu::GpuDevice;
    148
    149
    0
            let gpu = GpuDevice::new().map_err(TruenoError::InvalidInput)?;
    150
    151
    0
            gpu.convolve2d(
    152
    0
                self.as_slice(),
    153
    0
                kernel.as_slice(),
    154
    0
                result.data.as_mut_slice(),
    155
    0
                self.rows,
    156
    0
                self.cols,
    157
    0
                kernel.rows,
    158
    0
                kernel.cols,
    159
            )
    160
    0
            .map_err(TruenoError::InvalidInput)?;
    161
    162
    0
            Ok(result.clone())
    163
    0
        }
    164
    165
        /// Lookup embeddings by indices
    166
        ///
    167
        /// Performs embedding lookup where self is the embedding table with shape
    168
        /// `[vocab_size, embed_dim]` and indices specify which rows to select.
    169
        ///
    170
        /// # Arguments
    171
        ///
    172
        /// * `indices` - Slice of indices into the embedding table
    173
        ///
    174
        /// # Returns
    175
        ///
    176
        /// A matrix with shape `[indices.len(), embed_dim]` containing the selected rows
    177
        ///
    178
        /// # Errors
    179
        ///
    180
        /// Returns `InvalidInput` if any index is out of bounds
    181
        ///
    182
        /// # Example
    183
        ///
    184
        /// ```
    185
        /// use trueno::Matrix;
    186
        ///
    187
        /// // Create embedding table: 4 words, 3-dimensional embeddings
    188
        /// let embeddings = Matrix::from_vec(4, 3, vec![
    189
        ///     1.0, 2.0, 3.0,   // word 0
    190
        ///     4.0, 5.0, 6.0,   // word 1
    191
        ///     7.0, 8.0, 9.0,   // word 2
    192
        ///     10.0, 11.0, 12.0 // word 3
    193
        /// ]).unwrap();
    194
        ///
    195
        /// // Lookup embeddings for indices [1, 3, 0]
    196
        /// let result = embeddings.embedding_lookup(&[1, 3, 0]).unwrap();
    197
        ///
    198
        /// assert_eq!(result.rows(), 3);
    199
        /// assert_eq!(result.cols(), 3);
    200
        /// assert_eq!(result.get(0, 0), Some(&4.0)); // word 1
    201
        /// assert_eq!(result.get(1, 0), Some(&10.0)); // word 3
    202
        /// assert_eq!(result.get(2, 0), Some(&1.0)); // word 0
    203
        /// ```
    204
    0
        pub fn embedding_lookup(&self, indices: &[usize]) -> Result<Matrix<f32>, TruenoError> {
    205
            // Validate indices
    206
    0
            for (i, &idx) in indices.iter().enumerate() {
    207
    0
                if idx >= self.rows {
    208
    0
                    return Err(TruenoError::InvalidInput(format!(
    209
    0
                        "Index {} at position {} is out of bounds for embedding table with {} rows",
    210
    0
                        idx, i, self.rows
    211
    0
                    )));
    212
    0
                }
    213
            }
    214
    215
            // Handle empty indices
    216
    0
            if indices.is_empty() {
    217
    0
                return Ok(Matrix::zeros_with_backend(0, self.cols, self.backend));
    218
    0
            }
    219
    220
            // Allocate output matrix: [seq_len, embed_dim]
    221
    0
            let seq_len = indices.len();
    222
    0
            let embed_dim = self.cols;
    223
    0
            let mut result = Matrix::zeros_with_backend(seq_len, embed_dim, self.backend);
    224
    225
            // Copy rows from embedding table to result
    226
    0
            for (out_row, &idx) in indices.iter().enumerate() {
    227
    0
                let src_start = idx * embed_dim;
    228
    0
                let dst_start = out_row * embed_dim;
    229
    0
    230
    0
                result.data[dst_start..dst_start + embed_dim]
    231
    0
                    .copy_from_slice(&self.data[src_start..src_start + embed_dim]);
    232
    0
            }
    233
    234
    0
            Ok(result)
    235
    0
        }
    236
    237
        /// Lookup embeddings with gradient tracking support (for training)
    238
        ///
    239
        /// Returns both the embeddings and a sparse gradient accumulator.
    240
        /// This is useful for sparse gradient updates in training.
    241
        ///
    242
        /// # Arguments
    243
        ///
    244
        /// * `indices` - Slice of indices into the embedding table
    245
        ///
    246
        /// # Returns
    247
        ///
    248
        /// Tuple of (embeddings, unique_indices) where unique_indices can be used
    249
        /// for sparse gradient updates
    250
        ///
    251
        /// # Errors
    252
        ///
    253
        /// Returns `InvalidInput` if any index is out of bounds
    254
    0
        pub fn embedding_lookup_sparse(
    255
    0
            &self,
    256
    0
            indices: &[usize],
    257
    0
        ) -> Result<(Matrix<f32>, Vec<usize>), TruenoError> {
    258
    0
            let embeddings = self.embedding_lookup(indices)?;
    259
    260
            // Get unique indices for sparse gradient updates
    261
    0
            let mut unique: Vec<usize> = indices.to_vec();
    262
    0
            unique.sort_unstable();
    263
    0
            unique.dedup();
    264
    265
    0
            Ok((embeddings, unique))
    266
    0
        }
    267
    268
        /// 2D Max Pooling operation for CNN downsampling
    269
        ///
    270
        /// Applies max pooling over a 2D input tensor with specified kernel size and stride.
    271
        ///
    272
        /// # Arguments
    273
        /// * `kernel` - (kernel_height, kernel_width) pooling window size
    274
        /// * `stride` - (stride_height, stride_width) step size
    275
        ///
    276
        /// # Examples
    277
        /// ```
    278
        /// use trueno::matrix::Matrix;
    279
        /// let input = Matrix::from_vec(4, 4, vec![
    280
        ///     1.0, 2.0, 3.0, 4.0,
    281
        ///     5.0, 6.0, 7.0, 8.0,
    282
        ///     9.0, 10.0, 11.0, 12.0,
    283
        ///     13.0, 14.0, 15.0, 16.0,
    284
        /// ]).unwrap();
    285
        /// let pooled = input.max_pool2d((2, 2), (2, 2)).unwrap();
    286
        /// assert_eq!(pooled.shape(), (2, 2));
    287
        /// assert_eq!(pooled.get(0, 0), Some(&6.0));  // max of [1,2,5,6]
    288
        /// assert_eq!(pooled.get(1, 1), Some(&16.0)); // max of [11,12,15,16]
    289
        /// ```
    290
    0
        pub fn max_pool2d(
    291
    0
            &self,
    292
    0
            kernel: (usize, usize),
    293
    0
            stride: (usize, usize),
    294
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    295
    0
            let (kh, kw) = kernel;
    296
    0
            let (sh, sw) = stride;
    297
    298
    0
            if kh == 0 || kw == 0 || sh == 0 || sw == 0 {
    299
    0
                return Err(TruenoError::InvalidInput(
    300
    0
                    "Kernel and stride dimensions must be positive".into(),
    301
    0
                ));
    302
    0
            }
    303
    304
    0
            if kh > self.rows || kw > self.cols {
    305
    0
                return Err(TruenoError::InvalidInput(format!(
    306
    0
                    "Kernel size ({}, {}) larger than input ({}, {})",
    307
    0
                    kh, kw, self.rows, self.cols
    308
    0
                )));
    309
    0
            }
    310
    311
    0
            let out_h = (self.rows - kh) / sh + 1;
    312
    0
            let out_w = (self.cols - kw) / sw + 1;
    313
    0
            let mut result = Matrix::new(out_h, out_w);
    314
    315
    0
            for i in 0..out_h {
    316
    0
                for j in 0..out_w {
    317
    0
                    let mut max_val = f32::NEG_INFINITY;
    318
    0
                    for ki in 0..kh {
    319
    0
                        for kj in 0..kw {
    320
    0
                            let val = self.data[(i * sh + ki) * self.cols + (j * sw + kj)];
    321
    0
                            max_val = max_val.max(val);
    322
    0
                        }
    323
                    }
    324
    0
                    result.data[i * out_w + j] = max_val;
    325
                }
    326
            }
    327
    328
    0
            Ok(result)
    329
    0
        }
    330
    331
        /// 2D Average Pooling operation for CNN downsampling
    332
        ///
    333
        /// Applies average pooling over a 2D input tensor with specified kernel size and stride.
    334
        ///
    335
        /// # Arguments
    336
        /// * `kernel` - (kernel_height, kernel_width) pooling window size
    337
        /// * `stride` - (stride_height, stride_width) step size
    338
        ///
    339
        /// # Examples
    340
        /// ```
    341
        /// use trueno::matrix::Matrix;
    342
        /// let input = Matrix::from_vec(4, 4, vec![
    343
        ///     1.0, 2.0, 3.0, 4.0,
    344
        ///     5.0, 6.0, 7.0, 8.0,
    345
        ///     9.0, 10.0, 11.0, 12.0,
    346
        ///     13.0, 14.0, 15.0, 16.0,
    347
        /// ]).unwrap();
    348
        /// let pooled = input.avg_pool2d((2, 2), (2, 2)).unwrap();
    349
        /// assert_eq!(pooled.shape(), (2, 2));
    350
        /// assert!((pooled.get(0, 0).unwrap() - 3.5).abs() < 1e-5);  // avg of [1,2,5,6]
    351
        /// ```
    352
    0
        pub fn avg_pool2d(
    353
    0
            &self,
    354
    0
            kernel: (usize, usize),
    355
    0
            stride: (usize, usize),
    356
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    357
    0
            let (kh, kw) = kernel;
    358
    0
            let (sh, sw) = stride;
    359
    360
    0
            if kh == 0 || kw == 0 || sh == 0 || sw == 0 {
    361
    0
                return Err(TruenoError::InvalidInput(
    362
    0
                    "Kernel and stride dimensions must be positive".into(),
    363
    0
                ));
    364
    0
            }
    365
    366
    0
            if kh > self.rows || kw > self.cols {
    367
    0
                return Err(TruenoError::InvalidInput(format!(
    368
    0
                    "Kernel size ({}, {}) larger than input ({}, {})",
    369
    0
                    kh, kw, self.rows, self.cols
    370
    0
                )));
    371
    0
            }
    372
    373
    0
            let out_h = (self.rows - kh) / sh + 1;
    374
    0
            let out_w = (self.cols - kw) / sw + 1;
    375
    0
            let kernel_size = (kh * kw) as f32;
    376
    0
            let mut result = Matrix::new(out_h, out_w);
    377
    378
    0
            for i in 0..out_h {
    379
    0
                for j in 0..out_w {
    380
    0
                    let mut sum = 0.0;
    381
    0
                    for ki in 0..kh {
    382
    0
                        for kj in 0..kw {
    383
    0
                            sum += self.data[(i * sh + ki) * self.cols + (j * sw + kj)];
    384
    0
                        }
    385
                    }
    386
    0
                    result.data[i * out_w + j] = sum / kernel_size;
    387
                }
    388
            }
    389
    390
    0
            Ok(result)
    391
    0
        }
    392
    393
        /// Top-K selection: returns the k largest elements and their indices
    394
        ///
    395
        /// Useful for beam search, sampling, and ranking operations.
    396
        /// Searches row-major order and returns (values, indices) sorted descending.
    397
        ///
    398
        /// # Examples
    399
        /// ```
    400
        /// use trueno::matrix::Matrix;
    401
        /// let m = Matrix::from_vec(2, 3, vec![1.0, 5.0, 3.0, 2.0, 6.0, 4.0]).unwrap();
    402
        /// let (values, indices) = m.topk(2).unwrap();
    403
        /// assert_eq!(values, vec![6.0, 5.0]);
    404
        /// assert_eq!(indices, vec![4, 1]);  // flat indices
    405
        /// ```
    406
    0
        pub fn topk(&self, k: usize) -> Result<(Vec<f32>, Vec<usize>), TruenoError> {
    407
    0
            if k == 0 {
    408
    0
                return Ok((vec![], vec![]));
    409
    0
            }
    410
    411
    0
            let k = k.min(self.data.len());
    412
    0
            let mut indexed: Vec<(usize, f32)> = self.data.iter().copied().enumerate().collect();
    413
    414
            // Partial sort - only sort k elements
    415
    0
            indexed.select_nth_unstable_by(k.saturating_sub(1), |a, b| {
    416
    0
                b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
    417
    0
            });
    418
    419
    0
            indexed.truncate(k);
    420
    0
            indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    421
    422
    0
            let values: Vec<f32> = indexed.iter().map(|(_, v)| *v).collect();
    423
    0
            let indices: Vec<usize> = indexed.iter().map(|(i, _)| *i).collect();
    424
    425
    0
            Ok((values, indices))
    426
    0
        }
    427
    428
        /// Gather elements along axis using indices
    429
        ///
    430
        /// For 2D matrix with axis=0: output[i] = self[indices[i], :]
    431
        /// For 2D matrix with axis=1: output[:, i] = self[:, indices[i]]
    432
        ///
    433
        /// # Examples
    434
        /// ```
    435
        /// use trueno::matrix::Matrix;
    436
        /// let m = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    437
        /// let gathered = m.gather(&[2, 0], 0).unwrap();  // Select rows 2 and 0
    438
        /// assert_eq!(gathered.shape(), (2, 2));
    439
        /// assert_eq!(gathered.get(0, 0), Some(&5.0));  // Row 2
    440
        /// assert_eq!(gathered.get(1, 0), Some(&1.0));  // Row 0
    441
        /// ```
    442
    0
        pub fn gather(&self, indices: &[usize], axis: usize) -> Result<Matrix<f32>, TruenoError> {
    443
    0
            match axis {
    444
                0 => {
    445
                    // Gather rows
    446
    0
                    let mut result = Matrix::new(indices.len(), self.cols);
    447
    0
                    for (out_i, &idx) in indices.iter().enumerate() {
    448
    0
                        if idx >= self.rows {
    449
    0
                            return Err(TruenoError::InvalidInput(format!(
    450
    0
                                "Index {} out of bounds for axis 0 with size {}",
    451
    0
                                idx, self.rows
    452
    0
                            )));
    453
    0
                        }
    454
    0
                        for j in 0..self.cols {
    455
    0
                            result.data[out_i * self.cols + j] = self.data[idx * self.cols + j];
    456
    0
                        }
    457
                    }
    458
    0
                    Ok(result)
    459
                }
    460
                1 => {
    461
                    // Gather columns
    462
    0
                    let mut result = Matrix::new(self.rows, indices.len());
    463
    0
                    for i in 0..self.rows {
    464
    0
                        for (out_j, &idx) in indices.iter().enumerate() {
    465
    0
                            if idx >= self.cols {
    466
    0
                                return Err(TruenoError::InvalidInput(format!(
    467
    0
                                    "Index {} out of bounds for axis 1 with size {}",
    468
    0
                                    idx, self.cols
    469
    0
                                )));
    470
    0
                            }
    471
    0
                            result.data[i * indices.len() + out_j] = self.data[i * self.cols + idx];
    472
                        }
    473
                    }
    474
    0
                    Ok(result)
    475
                }
    476
    0
                _ => Err(TruenoError::InvalidInput(format!(
    477
    0
                    "Axis {} not supported for 2D matrix (use 0 or 1)",
    478
    0
                    axis
    479
    0
                ))),
    480
            }
    481
    0
        }
    482
    483
        /// Pad matrix with a constant value
    484
        ///
    485
        /// # Arguments
    486
        /// * `padding` - ((top, bottom), (left, right)) padding amounts
    487
        /// * `value` - constant value to pad with (usually 0.0)
    488
        ///
    489
        /// # Examples
    490
        /// ```
    491
        /// use trueno::matrix::Matrix;
    492
        /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    493
        /// let padded = m.pad(((1, 1), (1, 1)), 0.0).unwrap();
    494
        /// assert_eq!(padded.shape(), (4, 4));
    495
        /// assert_eq!(padded.get(0, 0), Some(&0.0));  // top-left padding
    496
        /// assert_eq!(padded.get(1, 1), Some(&1.0));  // original (0,0)
    497
        /// ```
    498
    0
        pub fn pad(
    499
    0
            &self,
    500
    0
            padding: ((usize, usize), (usize, usize)),
    501
    0
            value: f32,
    502
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    503
    0
            let ((top, bottom), (left, right)) = padding;
    504
    0
            let new_rows = self.rows + top + bottom;
    505
    0
            let new_cols = self.cols + left + right;
    506
    507
    0
            let mut result = Matrix::from_vec(new_rows, new_cols, vec![value; new_rows * new_cols])?;
    508
    509
            // Copy original data
    510
    0
            for i in 0..self.rows {
    511
    0
                for j in 0..self.cols {
    512
    0
                    result.data[(i + top) * new_cols + (j + left)] = self.data[i * self.cols + j];
    513
    0
                }
    514
            }
    515
    516
    0
            Ok(result)
    517
    0
        }
    518
    }
    519
    520
    #[cfg(test)]
    521
    mod tests {
    522
        use super::*;
    523
    524
        #[test]
    525
        fn test_convolve2d_basic() {
    526
            let input = Matrix::from_vec(3, 3, vec![1.0; 9]).unwrap();
    527
            let kernel = Matrix::from_vec(2, 2, vec![1.0; 4]).unwrap();
    528
            let result = input.convolve2d(&kernel).unwrap();
    529
            assert_eq!(result.rows(), 2);
    530
            assert_eq!(result.cols(), 2);
    531
            assert_eq!(result.get(0, 0), Some(&4.0));
    532
        }
    533
    534
        #[test]
    535
        fn test_embedding_lookup() {
    536
            let embeddings = Matrix::from_vec(4, 3, vec![
    537
                1.0, 2.0, 3.0,
    538
                4.0, 5.0, 6.0,
    539
                7.0, 8.0, 9.0,
    540
                10.0, 11.0, 12.0,
    541
            ]).unwrap();
    542
            let result = embeddings.embedding_lookup(&[1, 3]).unwrap();
    543
            assert_eq!(result.rows(), 2);
    544
            assert_eq!(result.get(0, 0), Some(&4.0));
    545
            assert_eq!(result.get(1, 0), Some(&10.0));
    546
        }
    547
    548
        #[test]
    549
        fn test_embedding_lookup_out_of_bounds() {
    550
            let embeddings = Matrix::from_vec(4, 3, vec![0.0; 12]).unwrap();
    551
            assert!(embeddings.embedding_lookup(&[5]).is_err());
    552
        }
    553
    554
        #[test]
    555
        fn test_max_pool2d() {
    556
            let input = Matrix::from_vec(4, 4, vec![
    557
                1.0, 2.0, 3.0, 4.0,
    558
                5.0, 6.0, 7.0, 8.0,
    559
                9.0, 10.0, 11.0, 12.0,
    560
                13.0, 14.0, 15.0, 16.0,
    561
            ]).unwrap();
    562
            let pooled = input.max_pool2d((2, 2), (2, 2)).unwrap();
    563
            assert_eq!(pooled.shape(), (2, 2));
    564
            assert_eq!(pooled.get(0, 0), Some(&6.0));
    565
            assert_eq!(pooled.get(1, 1), Some(&16.0));
    566
        }
    567
    568
        #[test]
    569
        fn test_avg_pool2d() {
    570
            let input = Matrix::from_vec(4, 4, vec![
    571
                1.0, 2.0, 3.0, 4.0,
    572
                5.0, 6.0, 7.0, 8.0,
    573
                9.0, 10.0, 11.0, 12.0,
    574
                13.0, 14.0, 15.0, 16.0,
    575
            ]).unwrap();
    576
            let pooled = input.avg_pool2d((2, 2), (2, 2)).unwrap();
    577
            assert_eq!(pooled.shape(), (2, 2));
    578
            assert!((pooled.get(0, 0).unwrap() - 3.5).abs() < 1e-5);
    579
        }
    580
    581
        #[test]
    582
        fn test_topk() {
    583
            let m = Matrix::from_vec(2, 3, vec![1.0, 5.0, 3.0, 2.0, 6.0, 4.0]).unwrap();
    584
            let (values, indices) = m.topk(2).unwrap();
    585
            assert_eq!(values, vec![6.0, 5.0]);
    586
            assert_eq!(indices, vec![4, 1]);
    587
        }
    588
    589
        #[test]
    590
        fn test_gather_rows() {
    591
            let m = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    592
            let gathered = m.gather(&[2, 0], 0).unwrap();
    593
            assert_eq!(gathered.shape(), (2, 2));
    594
            assert_eq!(gathered.get(0, 0), Some(&5.0));
    595
            assert_eq!(gathered.get(1, 0), Some(&1.0));
    596
        }
    597
    598
        #[test]
    599
        fn test_pad() {
    600
            let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    601
            let padded = m.pad(((1, 1), (1, 1)), 0.0).unwrap();
    602
            assert_eq!(padded.shape(), (4, 4));
    603
            assert_eq!(padded.get(0, 0), Some(&0.0));
    604
            assert_eq!(padded.get(1, 1), Some(&1.0));
    605
        }
    606
    }