Improve comments

sonos · Dec 11, 2024 · e78057f · e78057f
1 parent 4f43dfc
commit e78057f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/metal/src/ops/scaled_masked_softmax.rs b/metal/src/ops/scaled_masked_softmax.rs
@@ -5,8 +5,8 @@ use crate::MetalContext;
 use derive_new::new;
 use tract_core::internal::*;
 
-/// A = SOFTMAX(INPUT * SCALE + MASK, AXIS=1)
-/// Only input of rank of 3 is supported and softmax axis = 2
+/// A = SOFTMAX(INPUT * SCALE + MASK, AXIS=2)
+/// Only input of rank of 3 is supported
 #[derive(Clone, Debug, new, Hash)]
 pub struct MetalScaledMaskedSoftmax {
     pub scale: Arc<Tensor>,

diff --git a/metal/src/rewrite_rules/scaled_masked_softmax.rs b/metal/src/rewrite_rules/scaled_masked_softmax.rs
@@ -8,8 +8,8 @@ use tract_core::ops::binary::BinMiniOp;
 use tract_core::ops::math::{Add, Mul};
 use tract_core::ops::nn::{Softmax, SoftmaxExp};
 
-/// A = SOFTMAX(INPUT * SCALE + MASK, AXIS=1)
-/// Only input of rank of 3 is supported with softmax axis = 2
+/// A = SOFTMAX(INPUT * SCALE + MASK, AXIS=2)
+/// Only input of rank of 3 is supported.
 #[derive(Clone, Debug, Hash)]
 pub struct BasicScaledMaskedSoftmax {
     pub scale: Arc<Tensor>,
@@ -56,7 +56,7 @@ impl TypedOp for BasicScaledMaskedSoftmax {
     as_op!();
 }
 
-/// Search pattern => A = SOFTMAX(A * SCALE + MASK)
+/// Search pattern => A = SOFTMAX(A * SCALE + MASK, AXIS=2)
 pub fn as_scaled_masked_softmax_rule(
     _ctx: &(),
     model: &TypedModel,