#7826: Add doc for abs

tenstorrent · Sep 26, 2024 · 84678ff · 84678ff
1 parent c57584a
commit 84678ff
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 3 deletions.
diff --git a/tests/ttnn/sweep_tests/sweeps/sweeps/abs.py b/tests/ttnn/sweep_tests/sweeps/sweeps/abs.py
@@ -16,13 +16,27 @@
     "batch_sizes": [(1,)],
     "height": [384, 1024],
     "width": [1024, 4096],
-    "input_dtype": [ttnn.bfloat16],
+    "input_dtype": [ttnn.bfloat16, ttnn.bfloat8_b],
     "input_memory_config": [ttnn.DRAM_MEMORY_CONFIG],
     "output_memory_config": [ttnn.DRAM_MEMORY_CONFIG],
-    "layout": [ttnn.TILE_LAYOUT],
+    "layout": [ttnn.TILE_LAYOUT, ttnn.ROW_MAJOR_LAYOUT],
 }
 
 
+def skip(
+    batch_sizes,
+    height,
+    width,
+    input_dtype,
+    input_memory_config,
+    output_memory_config,
+    layout,
+) -> Tuple[bool, Optional[str]]:
+    if layout == ttnn.ROW_MAJOR_LAYOUT:
+        return True, "Skipped as ROW_MAJOR_LAYOUT not supported"
+    return False, None
+
+
 def run(
     batch_sizes,
     height,

diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp
@@ -1318,7 +1318,14 @@ void bind_dropout(py::module& module, const unary_operation_t& operation) {
 }  // namespace detail
 
 void py_module(py::module& module) {
-    detail::bind_unary_operation_overload_complex(module, ttnn::abs);
+    detail::bind_unary_operation_overload_complex(module, ttnn::abs,
+    R"doc(Supported dtypes, layouts, and ranks:
+
+        +----------------------------+---------------------------------+-------------------+
+        |     Dtypes                 |         Layouts                 |     Ranks         |
+        +----------------------------+---------------------------------+-------------------+
+        |    BFLOAT16, BFLOAT8_B     |          TILE                   |      2, 3, 4      |
+        +----------------------------+---------------------------------+-------------------+)doc");
     detail::bind_unary_operation(module, ttnn::acos);
     detail::bind_unary_operation(module, ttnn::asin);
     detail::bind_unary_operation(module, ttnn::atan);