diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp b/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp index f6e72d7edf84..84799557b36a 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp @@ -81,13 +81,13 @@ void bind_primitive_binary_operation(py::module& module, const binary_operation_ template -void bind_binary_operation(py::module& module, const binary_operation_t& operation, const std::string& description) { +void bind_binary_operation(py::module& module, const binary_operation_t& operation, const std::string& description, const std::string& math) { auto doc = fmt::format( R"doc( {2} .. math:: - \mathrm{{output\_tensor}} = {0}(\mathrm{{input\_tensor\_a,input\_tensor\_b}}). + {3} Args: input_tensor_a (ttnn.Tensor): the input tensor. @@ -113,7 +113,8 @@ void bind_binary_operation(py::module& module, const binary_operation_t& operati )doc", operation.base_name(), operation.python_fully_qualified_name(), - description); + description, + math); bind_registered_operation( module, @@ -328,7 +329,7 @@ void bind_div_like_ops(py::module& module, const binary_operation_t& operation, {2} .. math:: - \mathrm{{output\_tensor}} = {0}(\mathrm{{input\_tensor\_a,input\_tensor\_b}}). + \mathrm{{output\_tensor}} = \verb|{0}|(\mathrm{{input\_tensor\_a,input\_tensor\_b}}). Args: input_tensor_a (ttnn.Tensor): the input tensor. @@ -563,11 +564,14 @@ void bind_binary_overload_operation(py::module& module, const binary_operation_t } template -void bind_inplace_operation(py::module& module, const binary_operation_t& operation, const std::string& description) { +void bind_inplace_operation(py::module& module, const binary_operation_t& operation, const std::string& description, const std::string& math) { auto doc = fmt::format( R"doc( {2} + .. math:: + {3} + Args: input_tensor_a (ttnn.Tensor): the input tensor. input_tensor_b (ttnn.Tensor or Number): the input tensor. @@ -581,7 +585,8 @@ void bind_inplace_operation(py::module& module, const binary_operation_t& operat )doc", operation.base_name(), operation.python_fully_qualified_name(), - description); + description, + math); bind_registered_operation( module, @@ -614,7 +619,7 @@ void bind_logical_inplace_operation(py::module& module, const binary_operation_t {2} .. math:: - \mathrm{{output\_tensor}} = {0}(\mathrm{{input\_tensor\_a,input\_tensor\_b}}). + \mathrm{{output\_tensor}} = \verb|{0}|(\mathrm{{input\_tensor\_a,input\_tensor\_b}}). Args: input_tensor_a (ttnn.Tensor): the input tensor. @@ -706,7 +711,8 @@ void py_module(py::module& module) { detail::bind_binary_operation( module, ttnn::add, - R"doc(Adds :attr:`input_tensor_a` to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Adds :attr:`input_tensor_a` to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i + \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_inplace_operation( module, @@ -717,7 +723,8 @@ void py_module(py::module& module) { detail::bind_binary_operation( module, ttnn::subtract, - R"doc(Subtracts :attr:`input_tensor_b` from :attr:`input_tensor_a` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Subtracts :attr:`input_tensor_b` from :attr:`input_tensor_a` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i - \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_inplace_operation( module, @@ -728,7 +735,8 @@ void py_module(py::module& module) { detail::bind_binary_operation( module, ttnn::multiply, - R"doc(Multiplies :attr:`input_tensor_a` by :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Multiplies :attr:`input_tensor_a` by :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i * \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_inplace_operation( module, @@ -739,72 +747,86 @@ void py_module(py::module& module) { detail::bind_binary_operation( module, ttnn::eq, - R"doc(Compares if :attr:`input_tensor_a` is equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compares if :attr:`input_tensor_a` is equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i == \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::ne, - R"doc(Compares if :attr:`input_tensor_a` is not equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compares if :attr:`input_tensor_a` is not equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i != \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::lt, - R"doc(Compares if :attr:`input_tensor_a` is less than :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compares if :attr:`input_tensor_a` is less than :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i < \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::le, - R"doc(MCompares if :attr:`input_tensor_a` is less than or equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(MCompares if :attr:`input_tensor_a` is less than or equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i <= \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::gt, - R"doc(Compares if :attr:`input_tensor_a` is greater than :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compares if :attr:`input_tensor_a` is greater than :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i > \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::ge, - R"doc(Compares if :attr:`input_tensor_a` is greater than or equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compares if :attr:`input_tensor_a` is greater than or equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i >= \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::logical_and, - R"doc(Compute logical AND of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute logical AND of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i \& \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::logical_or, - R"doc(Compute logical OR of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute logical OR of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i | \mathrm{{input\_tensor\_b}}_i))doc"); detail::bind_binary_operation( module, ttnn::ldexp, - R"doc(Compute ldexp of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute ldexp of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}} = ldexp(\mathrm{{input\_tensor\_a,input\_tensor\_b}}))doc"); detail::bind_binary_operation( module, ttnn::logaddexp, - R"doc(Compute logaddexp of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute logaddexp of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}} = logaddexp(\mathrm{{input\_tensor\_a,input\_tensor\_b}}))doc"); detail::bind_binary_operation( module, ttnn::logaddexp2, - R"doc(Compute logaddexp2 of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute logaddexp2 of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}} = logaddexp2(\mathrm{{input\_tensor\_a,input\_tensor\_b}}))doc"); detail::bind_binary_operation( module, ttnn::squared_difference, - R"doc(Compute squared difference of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute squared difference of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}} = \verb|squared_difference|(\mathrm{{input\_tensor\_a,input\_tensor\_b}}))doc"); detail::bind_binary_operation( module, ttnn::bias_gelu, - R"doc(Compute bias_gelu of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Compute bias_gelu of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}} = \verb|bias_gelu|(\mathrm{{input\_tensor\_a,input\_tensor\_b}}))doc"); detail::bind_binary_operation( module, ttnn::divide, - R"doc(Divides :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc"); + R"doc(Divides :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a`)doc", + R"doc(\mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor\_a}}_i / \mathrm{{input\_tensor\_b}}_i))doc"); auto prim_module = module.def_submodule("prim", "Primitive binary operations"); @@ -968,32 +990,38 @@ void py_module(py::module& module) { detail::bind_inplace_operation( module, ttnn::gt_, - R"doc(Perform Greater than in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc"); + R"doc(Perform Greater than in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", + R"doc(\mathrm{{output\_tensor}} = (\mathrm{{input\_tensor\_a}} > \mathrm{{input\_tensor\_b}}))doc"); detail::bind_inplace_operation( module, ttnn::ge_, - R"doc(Perform Greater than or equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc"); + R"doc(Perform Greater than or equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", + R"doc(\mathrm{{output\_tensor}} = (\mathrm{{input\_tensor\_a}} >= \mathrm{{input\_tensor\_b}}))doc"); detail::bind_inplace_operation( module, ttnn::lt_, - R"doc(Perform Less than in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc"); + R"doc(Perform Less than in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", + R"doc(\mathrm{{output\_tensor}} = (\mathrm{{input\_tensor\_a}} < \mathrm{{input\_tensor\_b}}))doc"); detail::bind_inplace_operation( module, ttnn::le_, - R"doc(Perform Less than or equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc"); + R"doc(Perform Less than or equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", + R"doc(\mathrm{{output\_tensor}} = (\mathrm{{input\_tensor\_a}} <= \mathrm{{input\_tensor\_b}}))doc"); detail::bind_inplace_operation( module, ttnn::eq_, - R"doc(Perform Equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc"); + R"doc(Perform Equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", + R"doc(\mathrm{{output\_tensor}} = (\mathrm{{input\_tensor\_a}} == \mathrm{{input\_tensor\_b}}))doc"); detail::bind_inplace_operation( module, ttnn::ne_, - R"doc(Perform Not equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc"); + R"doc(Perform Not equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", + R"doc(\mathrm{{output\_tensor}} = (\mathrm{{input\_tensor\_a}} != \mathrm{{input\_tensor\_b}}))doc"); } diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp index 07a037938287..843ce85251e4 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp @@ -259,7 +259,7 @@ void bind_unary_operation_with_float_parameter( {4} .. math:: - \mathrm{{output\_tensor}}_i = {0}(\mathrm{{input\_tensor}}_i) + \mathrm{{output\_tensor}}_i = \verb|{0}|(\mathrm{{input\_tensor}}_i) Args: input_tensor (ttnn.Tensor): the input tensor. @@ -319,7 +319,7 @@ void bind_unary_operation_with_integer_parameter( {4} .. math:: - \mathrm{{output\_tensor}}_i = {0}(\mathrm{{input\_tensor}}_i) + \mathrm{{output\_tensor}}_i = \verb|{0}|(\mathrm{{input\_tensor}}_i) Args: input_tensor (ttnn.Tensor): the input tensor. @@ -535,7 +535,7 @@ void bind_sigmoid_accurate(py::module& module, const unary_operation_t& operatio Applies {0} to :attr:`input_tensor` element-wise. .. math:: - \mathrm{{output\_tensor}}_i = {0}(\mathrm{{input\_tensor}}_i) + \mathrm{{output\_tensor}}_i = \verb|{0}|(\mathrm{{input\_tensor}}_i) Args: input_tensor (ttnn.Tensor): the input tensor. @@ -1372,7 +1372,7 @@ void py_module(py::module& module) { detail::bind_unary_operation(module, ttnn::square, R"doc(\mathrm{{output\_tensor}}_i = square(\mathrm{{input\_tensor}}_i))doc"); detail::bind_unary_operation(module, ttnn::tan, R"doc(\mathrm{{output\_tensor}}_i = tan(\mathrm{{input\_tensor}}_i))doc"); detail::bind_unary_operation(module, ttnn::tanh, R"doc(\mathrm{{output\_tensor}}_i = tanh(\mathrm{{input\_tensor}}_i))doc"); - detail::bind_unary_operation(module, ttnn::log_sigmoid, R"doc(\mathrm{{output\_tensor}}_i = log_sigmoid(\mathrm{{input\_tensor}}_i))doc"); + detail::bind_unary_operation(module, ttnn::log_sigmoid, R"doc(\mathrm{{output\_tensor}}_i = \verb|log_sigmoid|(\mathrm{{input\_tensor}}_i))doc"); // Unaries with fast_and_approximate_mode detail::bind_unary_operation_with_fast_and_approximate_mode(module, ttnn::exp,