add tests and fix impl

Author: camfairchild

pallets/subtensor/src/epoch/math.rs

@@ -569,6 +569,33 @@ pub fn inplace_mask_diag(matrix: &mut [Vec<I32F32>]) {
     });
 }
 
+// Mask out the diagonal of the input matrix in-place, except for the diagonal entry at except_index.
+#[allow(dead_code)]
+pub fn inplace_mask_diag_except_index(matrix: &mut [Vec<I32F32>], except_index: u16) {
+    let Some(first_row) = matrix.first() else {
+        return;
+    };
+    if first_row.is_empty() {
+        return;
+    }
+    assert_eq!(matrix.len(), first_row.len());
+
+    let diag_at_index = matrix
+        .get(except_index as usize)
+        .and_then(|row| row.get(except_index as usize))
+        .cloned();
+
+    inplace_mask_diag(matrix);
+
+    matrix.get_mut(except_index as usize).map(|row| {
+        row.get_mut(except_index as usize).map(|value| {
+            if let Some(diag_at_index) = diag_at_index {
+                *value = diag_at_index;
+            }
+        })
+    });
+}
+
 // Return a new sparse matrix that replaces masked rows with an empty vector placeholder.
 #[allow(dead_code)]
 pub fn mask_rows_sparse(
@@ -611,23 +638,20 @@ pub fn mask_diag_sparse_except_index(
     sparse_matrix: &[Vec<(u16, I32F32)>],
     except_index: u16,
 ) -> Vec<Vec<(u16, I32F32)>> {
-    // Store the diagonal entry at except_index
-    let diag_at_index = sparse_matrix
-        .get(except_index as usize)
-        .and_then(|row| row.get(except_index as usize))
-        .cloned();
-    // Mask out the diagonal
-    let mut result = mask_diag_sparse(sparse_matrix);
-    // Replace the diagonal entry at except_index using only get_mut or map
-    result.get_mut(except_index as usize).map(|row| {
-        row.get_mut(except_index as usize).map(|value| {
-            if let Some(diag_at_index) = diag_at_index {
-                *value = diag_at_index;
-            }
+    sparse_matrix
+        .iter()
+        .enumerate()
+        .map(|(i, sparse_row)| {
+            sparse_row
+                .iter()
+                .filter(|(j, _)| {
+                    // Is not a diagonal OR is the diagonal at except_index
+                    i != (*j as usize) || (i == except_index as usize && *j == except_index)
+                })
+                .copied()
+                .collect()
         })
-    });
-
-    result
+        .collect()
 }
 
 // Remove cells from sparse matrix where the mask function of two vectors is true.

pallets/subtensor/src/epoch/run_epoch.rs

@@ -111,6 +111,9 @@ impl<T: Config> Pallet<T> {
         // == Weights ==
         // =============
 
+        // Get owner uid.
+        let owner_uid: Option<u16> = Self::get_owner_uid(netuid);
+
         // Access network weights row unnormalized.
         let mut weights: Vec<Vec<I32F32>> = Self::get_weights(netuid);
         log::trace!("W:\n{:?}\n", &weights);
@@ -119,7 +122,13 @@ impl<T: Config> Pallet<T> {
         inplace_mask_rows(&validator_forbids, &mut weights);
         log::trace!("W (permit): {:?}", &weights);
 
-        // Remove self-weight by masking diagonal.
+        // Remove self-weight by masking diagonal; keep owner_uid self-weight.
+        if let Some(owner_uid) = owner_uid {
+            inplace_mask_diag_except_index(&mut weights, owner_uid);
+        } else {
+            inplace_mask_diag(&mut weights);
+        }
+
         inplace_mask_diag(&mut weights);
         log::trace!("W (permit+diag):\n{:?}\n", &weights);
 

pallets/subtensor/src/tests/math.rs

@@ -58,7 +58,13 @@ fn assert_sparse_mat_compare(
 ) {
     assert!(ma.len() == mb.len());
     for row in 0..ma.len() {
-        assert!(ma[row].len() == mb[row].len());
+        assert!(
+            ma[row].len() == mb[row].len(),
+            "row: {}, ma: {:?}, mb: {:?}",
+            row,
+            ma[row],
+            mb[row]
+        );
         for j in 0..ma[row].len() {
             assert!(ma[row][j].0 == mb[row][j].0); // u16
             assert_float_compare(ma[row][j].1, mb[row][j].1, epsilon) // I32F32
@@ -1034,6 +1040,27 @@ fn test_math_inplace_mask_diag() {
     );
 }
 
+#[test]
+fn test_math_inplace_mask_diag_except_index() {
+    let vector: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8., 9.];
+    let rows = 3;
+
+    for i in 0..rows {
+        let mut target: Vec<f32> = vec![0., 2., 3., 4., 0., 6., 7., 8., 0.];
+        let row = i * rows;
+        let col = i;
+        target[row + col] = vector[row + col];
+
+        let mut mat = vec_to_mat_fixed(&vector, rows, false);
+        inplace_mask_diag_except_index(&mut mat, i as u16);
+        assert_mat_compare(
+            &mat,
+            &vec_to_mat_fixed(&target, rows, false),
+            I32F32::from_num(0),
+        );
+    }
+}
+
 #[test]
 fn test_math_mask_rows_sparse() {
     let input: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8., 9.];
@@ -1105,6 +1132,58 @@ fn test_math_mask_diag_sparse() {
     );
 }
 
+#[test]
+fn test_math_mask_diag_sparse_except_index() {
+    let rows = 3;
+
+    let vector: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8., 9.];
+    let mat = vec_to_sparse_mat_fixed(&vector, rows, false);
+
+    for i in 0..rows {
+        let mut target: Vec<f32> = vec![0., 2., 3., 4., 0., 6., 7., 8., 0.];
+        let row = i * rows;
+        let col = i;
+        target[row + col] = vector[row + col];
+
+        let result = mask_diag_sparse_except_index(&mat, i as u16);
+        let target_as_mat = vec_to_sparse_mat_fixed(&target, rows, false);
+
+        assert_sparse_mat_compare(&result, &target_as_mat, I32F32::from_num(0));
+    }
+
+    let vector: Vec<f32> = vec![1., 0., 0., 0., 5., 0., 0., 0., 9.];
+    let mat = vec_to_sparse_mat_fixed(&vector, rows, false);
+
+    for i in 0..rows {
+        let mut target: Vec<f32> = vec![0., 0., 0., 0., 0., 0., 0., 0., 0.];
+        let row = i * rows;
+        let col = i;
+        target[row + col] = vector[row + col];
+
+        let result = mask_diag_sparse_except_index(&mat, i as u16);
+        let target_as_mat = vec_to_sparse_mat_fixed(&target, rows, false);
+        assert_eq!(result.len(), target_as_mat.len());
+
+        assert_sparse_mat_compare(&result, &target_as_mat, I32F32::from_num(0));
+    }
+
+    for i in 0..rows {
+        let vector: Vec<f32> = vec![0., 0., 0., 0., 0., 0., 0., 0., 0.];
+        let mat = vec_to_sparse_mat_fixed(&vector, rows, false);
+
+        let mut target: Vec<f32> = vec![0., 0., 0., 0., 0., 0., 0., 0., 0.];
+        let row = i * rows;
+        let col = i;
+        target[row + col] = vector[row + col];
+
+        let result = mask_diag_sparse_except_index(&mat, i as u16);
+        let target_as_mat = vec_to_sparse_mat_fixed(&target, rows, false);
+        assert_eq!(result.len(), target_as_mat.len());
+
+        assert_sparse_mat_compare(&result, &target_as_mat, I32F32::from_num(0));
+    }
+}
+
 #[test]
 fn test_math_vec_mask_sparse_matrix() {
     let vector: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8., 9.];