Add Buffer::is_square() and Buffer::is_vector() convenience predicates

crates/mohu-buffer/src/buffer.rs

@@ -557,6 +557,16 @@ impl Buffer {
     pub fn is_f_contiguous(&self)  -> bool { self.layout.is_f_contiguous() }
     /// Returns `true` if this buffer is contiguous in either C or F order.
     pub fn is_contiguous(&self)    -> bool { self.layout.is_contiguous() }
+    /// Returns true if the buffer is 0D (a scalar).
+    pub fn is_scalar_shape(&self) -> bool { self.ndim() == 0 }
+    /// Returns true if the buffer is 1D (a vector).
+    pub fn is_vector(&self) -> bool { self.ndim() == 1 }
+    /// Returns true if the buffer is 2D (a matrix).
+    pub fn is_matrix(&self) -> bool { self.ndim() == 2 }
+    /// Returns true if the buffer is 2D and both dimensions are equal.
+    pub fn is_square(&self) -> bool {
+        self.ndim() == 2 && self.shape()[0] == self.shape()[1]
+    }
     /// Returns `true` if the backing memory is SIMD-aligned.
     pub fn is_aligned(&self)       -> bool { self.flags.contains(BufferFlags::ALIGNED) }
     /// Returns `true` if the backing memory is shared with other `Buffer` instances.

crates/mohu-buffer/tests/integration.rs

@@ -48,6 +48,46 @@ fn from_slice_reshape_to_2d() {
     assert_eq!(buf.get::<i32>(&[1, 0]).unwrap(), 4);
 }
 
+#[test]
+fn shape_predicates_vector_matrix_scalar() {
+    let vector = Buffer::from_slice(&[1.0_f64, 2.0, 3.0]).unwrap();
+    assert!(vector.is_vector());
+    assert!(!vector.is_matrix());
+    assert!(!vector.is_square());
+    assert!(!vector.is_scalar_shape());
+
+    let matrix = Buffer::from_slice(&[1.0_f64, 2.0, 3.0, 4.0]).unwrap().reshape(&[2, 2]).unwrap();
+    assert!(matrix.is_matrix());
+    assert!(matrix.is_square());
+    assert!(!matrix.is_vector());
+    assert!(!matrix.is_scalar_shape());
+
+    let scalar = Buffer::zeros(DType::F64, &[]).unwrap();
+    assert!(scalar.is_scalar_shape());
+    assert!(!scalar.is_vector());
+    assert!(!scalar.is_matrix());
+    assert!(!scalar.is_square());
+}
+
+#[test]
+fn shape_predicates_square_and_rectangular() {
+    let square = Buffer::from_slice(&[1_i32]).unwrap().reshape(&[1, 1]).unwrap();
+    assert!(square.is_square());
+    assert!(square.is_matrix());
+
+    let rect = Buffer::from_slice(&[1_i32, 2, 3, 4, 5, 6]).unwrap().reshape(&[2, 3]).unwrap();
+    assert!(rect.is_matrix());
+    assert!(!rect.is_square());
+}
+
+#[test]
+fn shape_predicates_non_2d_false_for_square() {
+    let tensor = Buffer::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]).unwrap().reshape(&[2, 2, 1]).unwrap();
+    assert!(!tensor.is_square());
+    assert!(!tensor.is_matrix());
+    assert!(!tensor.is_vector());
+}
+
 #[test]
 fn get_set_roundtrip() {
     let mut buf = Buffer::zeros(DType::F64, &[4, 4]).unwrap();