documentation (#179)

sdpython · May 7, 2024 · e2e4f39 · e2e4f39
1 parent 36bb43f
commit e2e4f39
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 13 deletions.
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -10,7 +10,7 @@ Change Logs
 * :pr:`169`: adds custom operator ReplaceZero on CUDA
 * :pr:`168`: adds custom operator MulSigmoid on CUDA
 * :pr:`167`: adds custom operator Rotary on CUDA
-* :pr:`166`: adds custom operators AddMul, MulAdd on CUDA
+* :pr:`166`, :pr:`178`: adds custom operators AddMul, MulAdd on CUDA
 * :pr:`165`: adds custom operators AddAddAdd, MulMulMul on CUDA
 * :pr:`163`: use onnxruntime==1.17.3 as default
 * :pr:`162`: add ScatterNDOfShape implementation on CUDA without atomics

diff --git a/onnx_extended/ortops/optim/cuda/__init__.py b/onnx_extended/ortops/optim/cuda/__init__.py
@@ -24,10 +24,10 @@ def documentation() -> List[str]:
             textwrap.dedent,
             [
                 """
-    onnx_extended.ortops.optim.cuda.AddMul
+    onnx_extended.ortops.optim.cuda.AddAdd
     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-    Two consecutive element-wise Add, Mul assuming
+    Two consecutive element-wise addition assuming
     all tensors have the same shape
     (broadcast limited to the first dimensions).
 
@@ -43,17 +43,17 @@ def documentation() -> List[str]:
 
     **Outputs**
 
-    * (A+B)*C (T): element-wise
+    * A+B+C (T): element-wise
 
     **Constraints**
 
     * T: float, float16
     """,
                 """
-    onnx_extended.ortops.optim.cuda.AddAdd
-    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    onnx_extended.ortops.optim.cuda.AddAddAdd
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-    Two consecutive element-wise addition assuming
+    Three consecutive element-wise addition assuming
     all tensors have the same shape
     (broadcast limited to the first dimensions).
 
@@ -66,37 +66,41 @@ def documentation() -> List[str]:
     * A (T): tensor of type T
     * B (T): tensor of type T
     * C (T): tensor of type T
+    * D (T): tensor of type T
 
     **Outputs**
 
-    * A+B+C (T): element-wise
+    * A+B+C+D (T): element-wise
 
     **Constraints**
 
     * T: float, float16
     """,
                 """
-    onnx_extended.ortops.optim.cuda.AddAddAdd
-    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+    onnx_extended.ortops.optim.cuda.AddMul
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-    Three consecutive element-wise addition assuming
+    Two consecutive element-wise Add, Mul assuming
     all tensors have the same shape
     (broadcast limited to the first dimensions).
 
     **Provider**
     
     CUDAExecutionProvider
     
+    **Attributes**
+    
+    * transposeMiddle: bool, if True, applies transposition [0, 2, 1, 3] on the result
+    
     **Inputs**
     
     * A (T): tensor of type T
     * B (T): tensor of type T
     * C (T): tensor of type T
-    * D (T): tensor of type T
 
     **Outputs**
 
-    * A+B+C+D (T): element-wise
+    * (A+B)*C (T): element-wise
 
     **Constraints**
 
@@ -142,7 +146,11 @@ def documentation() -> List[str]:
     **Provider**
     
     CUDAExecutionProvider
+        
+    **Attributes**
     
+    * transposeMiddle: bool, if True, applies transposition [0, 2, 1, 3] on the result
+
     **Inputs**
     
     * A (T): tensor of type T