Fixed sparse attention docs warning

LongxingTan · Jun 13, 2023 · 2662e25 · 2662e25
1 parent 2f2f2a5
commit 2662e25
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/tfts/layers/attention_layer.py b/tfts/layers/attention_layer.py
@@ -239,7 +239,16 @@ def __init__(self, hidden_size: int, num_heads: int, attention_dropout: float =
     def build(self, input_shape):
         super().build(input_shape)
 
-    def call(self, x):
+    def call(self, x, mask=None):
+        """Sparse attention
+
+        Parameters
+        ----------
+        x : tf.Tensor
+            _description_
+        mask : tf.Tensor, optional
+            _description_, by default None
+        """
         return
 
     def get_config(self):
@@ -254,14 +263,14 @@ def __init__(self, **kwargs):
     def build(self, input_shape):
         super().build(input_shape)
 
-    def call(self, x, x_mask=None):
-        """_summary_
+    def call(self, x, mask=None):
+        """Fast attention
 
         Parameters
         ----------
-        x : _type_
+        x : tf.Tensor
             _description_
-        x_mask : _type_, optional
+        mask : _type_, optional
             _description_, by default None
         """
         return