review

Author: crutcher

crates/burn-tensor/src/tensor/api/base.rs

@@ -895,8 +895,8 @@ where
                 size,
             );
             assert!(
-                dim < size,
-                "Expected: dim < size: found dim={}, size={}",
+                dim < self.shape().num_dims(),
+                "Expected: dim < num_dims: found dim={}, num_dims={}",
                 dim,
                 size,
             );

crates/burn-tensor/src/tensor/indexing/mod.rs

@@ -108,7 +108,7 @@ where
 /// ## Arguments
 ///
 /// * `idx` - The dimension index to canonicalize.
-/// * `rank` - The size of the dimension.
+/// * `rank` - The number of dimensions.
 /// * `wrap_scalar` - If true, pretend scalars have rank=1.
 ///
 /// ## Returns