Skip to content

Commit

Permalink
#12757: update math function for ops with underscore
Browse files Browse the repository at this point in the history
  • Loading branch information
@bharane-ab
bharane-ab committed Sep 27, 2024
1 parent 9b9b524 commit 05ad63b
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 34 deletions.
88 changes: 58 additions & 30 deletions ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,13 +81,13 @@ void bind_primitive_binary_operation(py::module& module, const binary_operation_


template <typename binary_operation_t>
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.
Expand All @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -563,11 +564,14 @@ void bind_binary_overload_operation(py::module& module, const binary_operation_t
}

template <typename binary_operation_t>
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.
Expand All @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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");

Expand Down Expand Up @@ -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");

}

Expand Down
8 changes: 4 additions & 4 deletions ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 05ad63b

Please sign in to comment.