Dev scalar op

oneflow/api/python/functional/python_functions.cpp

@@ -339,6 +339,234 @@ py::object PyScatter(py::args py_args, py::kwargs py_kwargs) {
   return py::cast(result.GetPtrOrThrow());
 }
 
+py::object PyEqual(py::args py_args, py::kwargs py_kwargs) {
+  // "broadcast_equal(Tensor x, Tensor y)"
+  // "scalar_logical_equal(Tensor in, Scalar scalar)"
+  PyObject* args = py_args.ptr();
+  size_t nargs = PyTuple_Size(args);
+  CHECK_EQ_OR_THROW(nargs, 2) << "2 positional inputs are required.";
+  const auto& result = [&]() -> Maybe<Tensor> {  // NOLINT
+    PyObject* input = PyTuple_GetItem(args, 0);
+    PyObject* other = PyTuple_GetItem(args, 1);
+    bool input_is_tensor = PyTensorCheck(input);
+    bool other_is_tensor = PyTensorCheck(other);
+    CHECK_OR_RETURN(input_is_tensor || other_is_tensor) << "Inputs must have one tensor at least.";
+    CHECK_OR_RETURN(PyTensorCheck(input) || PyScalarCheck(input))
+        << "The first input should be a tensor or scalar.";
+    CHECK_OR_RETURN(PyTensorCheck(other) || PyScalarCheck(other))
+        << "The second input should be a tensor or scalar.";
+
+    if (PyTensorCheck(input) && PyTensorCheck(other)) {
+      auto a = JUST(PyUnpackTensor(input));
+      auto b = JUST(PyUnpackTensor(other));
+      return functional::BroadcastEqual(a, b);
+    } else {
+      if (PyTensorCheck(input)) {
+        CHECK_OR_RETURN(PyScalarCheck(other)) << "The second input should be a scalar.";
+        auto a = JUST(PyUnpackTensor(input));
+        auto b = *JUST(PyUnpackScalar(other));
+        return functional::ScalarLogicalEqual(a, b);
+      } else {
+        CHECK_OR_RETURN(PyScalarCheck(input)) << "The first input should be a scalar.";
+        auto a = *JUST(PyUnpackScalar(input));
+        auto b = JUST(PyUnpackTensor(other));
+        return functional::ScalarLogicalEqual(b, a);
+      }
+    }
+  }();
+  return py::cast(result.GetPtrOrThrow());
+}
+
+py::object PyNotEqual(py::args py_args, py::kwargs py_kwargs) {
+  // "broadcast_not_equal(Tensor x, Tensor y)"
+  // "scalar_logical_not_equal(Tensor in, Scalar scalar)"
+  PyObject* args = py_args.ptr();
+  size_t nargs = PyTuple_Size(args);
+  CHECK_EQ_OR_THROW(nargs, 2) << "2 positional inputs are required.";
+  const auto& result = [&]() -> Maybe<Tensor> {  // NOLINT
+    PyObject* input = PyTuple_GetItem(args, 0);
+    PyObject* other = PyTuple_GetItem(args, 1);
+    bool input_is_tensor = PyTensorCheck(input);
+    bool other_is_tensor = PyTensorCheck(other);
+    CHECK_OR_RETURN(input_is_tensor || other_is_tensor) << "Inputs must have one tensor at least.";
+    CHECK_OR_RETURN(PyTensorCheck(input) || PyScalarCheck(input))
+        << "The first input should be a tensor or scalar.";
+    CHECK_OR_RETURN(PyTensorCheck(other) || PyScalarCheck(other))
+        << "The second input should be a tensor or scalar.";
+
+    if (PyTensorCheck(input) && PyTensorCheck(other)) {
+      auto a = JUST(PyUnpackTensor(input));
+      auto b = JUST(PyUnpackTensor(other));
+      return functional::BroadcastNotEqual(a, b);
+    } else {
+      if (PyTensorCheck(input)) {
+        CHECK_OR_RETURN(PyScalarCheck(other)) << "The second input should be a scalar.";
+        auto a = JUST(PyUnpackTensor(input));
+        auto b = *JUST(PyUnpackScalar(other));
+        return functional::ScalarLogicalNotEqual(a, b);
+      } else {
+        CHECK_OR_RETURN(PyScalarCheck(input)) << "The first input should be a scalar.";
+        auto a = *JUST(PyUnpackScalar(input));
+        auto b = JUST(PyUnpackTensor(other));
+        return functional::ScalarLogicalNotEqual(b, a);
+      }
+    }
+  }();
+  return py::cast(result.GetPtrOrThrow());
+}
+
+py::object PyGreater(py::args py_args, py::kwargs py_kwargs) {
+  // "broadcast_greater(Tensor x, Tensor y)"
+  // "scalar_logical_greater(Tensor in, Scalar scalar)"
+  PyObject* args = py_args.ptr();
+  size_t nargs = PyTuple_Size(args);
+  CHECK_EQ_OR_THROW(nargs, 2) << "2 positional inputs are required.";
+  const auto& result = [&]() -> Maybe<Tensor> {  // NOLINT
+    PyObject* input = PyTuple_GetItem(args, 0);
+    PyObject* other = PyTuple_GetItem(args, 1);
+    bool input_is_tensor = PyTensorCheck(input);
+    bool other_is_tensor = PyTensorCheck(other);
+    CHECK_OR_RETURN(input_is_tensor || other_is_tensor) << "Inputs must have one tensor at least.";
+    CHECK_OR_RETURN(PyTensorCheck(input) || PyScalarCheck(input))
+        << "The first input should be a tensor or scalar.";
+    CHECK_OR_RETURN(PyTensorCheck(other) || PyScalarCheck(other))
+        << "The second input should be a tensor or scalar.";
+
+    if (PyTensorCheck(input) && PyTensorCheck(other)) {
+      auto a = JUST(PyUnpackTensor(input));
+      auto b = JUST(PyUnpackTensor(other));
+      return functional::BroadcastGreater(a, b);
+    } else {
+      if (PyTensorCheck(input)) {
+        CHECK_OR_RETURN(PyScalarCheck(other)) << "The second input should be a scalar.";
+        auto a = JUST(PyUnpackTensor(input));
+        auto b = *JUST(PyUnpackScalar(other));
+        return functional::ScalarLogicalGreater(a, b);
+      } else {
+        CHECK_OR_RETURN(PyScalarCheck(input)) << "The first input should be a scalar.";
+        auto a = *JUST(PyUnpackScalar(input));
+        auto b = JUST(PyUnpackTensor(other));
+        return functional::ScalarLogicalGreater(b, a);
+      }
+    }
+  }();
+  return py::cast(result.GetPtrOrThrow());
+}
+
+py::object PyGreaterEqual(py::args py_args, py::kwargs py_kwargs) {
+  // "broadcast_greater_equal(Tensor x, Tensor y)"
+  // "scalar_logical_greater_equal(Tensor in, Scalar scalar)"
+  PyObject* args = py_args.ptr();
+  size_t nargs = PyTuple_Size(args);
+  CHECK_EQ_OR_THROW(nargs, 2) << "2 positional inputs are required.";
+  const auto& result = [&]() -> Maybe<Tensor> {  // NOLINT
+    PyObject* input = PyTuple_GetItem(args, 0);
+    PyObject* other = PyTuple_GetItem(args, 1);
+    bool input_is_tensor = PyTensorCheck(input);
+    bool other_is_tensor = PyTensorCheck(other);
+    CHECK_OR_RETURN(input_is_tensor || other_is_tensor) << "Inputs must have one tensor at least.";
+    CHECK_OR_RETURN(PyTensorCheck(input) || PyScalarCheck(input))
+        << "The first input should be a tensor or scalar.";
+    CHECK_OR_RETURN(PyTensorCheck(other) || PyScalarCheck(other))
+        << "The second input should be a tensor or scalar.";
+
+    if (PyTensorCheck(input) && PyTensorCheck(other)) {
+      auto a = JUST(PyUnpackTensor(input));
+      auto b = JUST(PyUnpackTensor(other));
+      return functional::BroadcastGreaterEqual(a, b);
+    } else {
+      if (PyTensorCheck(input)) {
+        CHECK_OR_RETURN(PyScalarCheck(other)) << "The second input should be a scalar.";
+        auto a = JUST(PyUnpackTensor(input));
+        auto b = *JUST(PyUnpackScalar(other));
+        return functional::ScalarLogicalGreaterEqual(a, b);
+      } else {
+        CHECK_OR_RETURN(PyScalarCheck(input)) << "The first input should be a scalar.";
+        auto a = *JUST(PyUnpackScalar(input));
+        auto b = JUST(PyUnpackTensor(other));
+        return functional::ScalarLogicalGreaterEqual(b, a);
+      }
+    }
+  }();
+  return py::cast(result.GetPtrOrThrow());
+}
+
+py::object PyLess(py::args py_args, py::kwargs py_kwargs) {
+  // "broadcast_less(Tensor x, Tensor y)"
+  // "scalar_logical_less(Tensor in, Scalar scalar)"
+  PyObject* args = py_args.ptr();
+  size_t nargs = PyTuple_Size(args);
+  CHECK_EQ_OR_THROW(nargs, 2) << "2 positional inputs are required.";
+  const auto& result = [&]() -> Maybe<Tensor> {  // NOLINT
+    PyObject* input = PyTuple_GetItem(args, 0);
+    PyObject* other = PyTuple_GetItem(args, 1);
+    bool input_is_tensor = PyTensorCheck(input);
+    bool other_is_tensor = PyTensorCheck(other);
+    CHECK_OR_RETURN(input_is_tensor || other_is_tensor) << "Inputs must have one tensor at least.";
+    CHECK_OR_RETURN(PyTensorCheck(input) || PyScalarCheck(input))
+        << "The first input should be a tensor or scalar.";
+    CHECK_OR_RETURN(PyTensorCheck(other) || PyScalarCheck(other))
+        << "The second input should be a tensor or scalar.";
+
+    if (PyTensorCheck(input) && PyTensorCheck(other)) {
+      auto a = JUST(PyUnpackTensor(input));
+      auto b = JUST(PyUnpackTensor(other));
+      return functional::BroadcastLess(a, b);
+    } else {
+      if (PyTensorCheck(input)) {
+        CHECK_OR_RETURN(PyScalarCheck(other)) << "The second input should be a scalar.";
+        auto a = JUST(PyUnpackTensor(input));
+        auto b = *JUST(PyUnpackScalar(other));
+        return functional::ScalarLogicalLess(a, b);
+      } else {
+        CHECK_OR_RETURN(PyScalarCheck(input)) << "The first input should be a scalar.";
+        auto a = *JUST(PyUnpackScalar(input));
+        auto b = JUST(PyUnpackTensor(other));
+        return functional::ScalarLogicalLess(b, a);
+      }
+    }
+  }();
+  return py::cast(result.GetPtrOrThrow());
+}
+
+py::object PyLessEqual(py::args py_args, py::kwargs py_kwargs) {
+  // "broadcast_less_equal(Tensor x, Tensor y)"
+  // "scalar_logical_less_equal(Tensor in, Scalar scalar)"
+  PyObject* args = py_args.ptr();
+  size_t nargs = PyTuple_Size(args);
+  CHECK_EQ_OR_THROW(nargs, 2) << "2 positional inputs are required.";
+  const auto& result = [&]() -> Maybe<Tensor> {  // NOLINT
+    PyObject* input = PyTuple_GetItem(args, 0);
+    PyObject* other = PyTuple_GetItem(args, 1);
+    bool input_is_tensor = PyTensorCheck(input);
+    bool other_is_tensor = PyTensorCheck(other);
+    CHECK_OR_RETURN(input_is_tensor || other_is_tensor) << "Inputs must have one tensor at least.";
+    CHECK_OR_RETURN(PyTensorCheck(input) || PyScalarCheck(input))
+        << "The first input should be a tensor or scalar.";
+    CHECK_OR_RETURN(PyTensorCheck(other) || PyScalarCheck(other))
+        << "The second input should be a tensor or scalar.";
+
+    if (PyTensorCheck(input) && PyTensorCheck(other)) {
+      auto a = JUST(PyUnpackTensor(input));
+      auto b = JUST(PyUnpackTensor(other));
+      return functional::BroadcastLessEqual(a, b);
+    } else {
+      if (PyTensorCheck(input)) {
+        CHECK_OR_RETURN(PyScalarCheck(other)) << "The second input should be a scalar.";
+        auto a = JUST(PyUnpackTensor(input));
+        auto b = *JUST(PyUnpackScalar(other));
+        return functional::ScalarLogicalLessEqual(a, b);
+      } else {
+        CHECK_OR_RETURN(PyScalarCheck(input)) << "The first input should be a scalar.";
+        auto a = *JUST(PyUnpackScalar(input));
+        auto b = JUST(PyUnpackTensor(other));
+        return functional::ScalarLogicalLessEqual(b, a);
+      }
+    }
+  }();
+  return py::cast(result.GetPtrOrThrow());
+}
+
 }  // namespace functional
 }  // namespace one
 
@@ -352,6 +580,12 @@ ONEFLOW_API_PYBIND11_MODULE("F", m) {
   m.def("pow", &functional::PyPow);
   m.def("clamp", &functional::PyClamp);
   m.def("scatter", &functional::PyScatter);
+  m.def("equal", &functional::PyEqual);
+  m.def("not_equal", &functional::PyNotEqual);
+  m.def("greater", &functional::PyGreater);
+  m.def("greater_equal", &functional::PyGreaterEqual);
+  m.def("less", &functional::PyLess);
+  m.def("less_equal", &functional::PyLessEqual);
 }
 
 }  // namespace oneflow

oneflow/core/functional/functional_api.yaml

@@ -950,6 +950,32 @@
   signature: "Tensor ReduceSumLike(Tensor in, Tensor like, *,Int32List axis)"
   bind_python: True
 
+- name: "scalar_logical_equal"
+  signature: "Tensor ScalarLogicalEqual(Tensor in, Scalar scalar)"
+  bind_python: False
+
+- name: "scalar_logical_not_equal"
+  signature: "Tensor ScalarLogicalNotEqual(Tensor in, Scalar scalar)"
+  bind_python: False
+
+- name: "scalar_logical_greater"
+  signature: "Tensor ScalarLogicalGreater(Tensor in, Scalar scalar)"
+  bind_python: False
+
+- name: "scalar_logical_greater_equal"
+  signature: "Tensor ScalarLogicalGreaterEqual(Tensor in, Scalar scalar)"
+  bind_python: False
+
+- name: "scalar_logical_less"
+  signature: "Tensor ScalarLogicalLess(Tensor in, Scalar scalar)"
+  bind_python: False
+
+- name: "scalar_logical_less_equal"
+  signature: "Tensor ScalarLogicalLessEqual(Tensor in, Scalar scalar)"
+  bind_python: False
+
+- name: "split"
+  signature: "TensorTuple Split(Tensor x, *, Int64 split_size, Int64 dim=0)"
 - name: "rand"
   signature: "Tensor Rand(*, Shape shape, DataType dtype=None, Device device=None, Generator generator=None)"
   bind_python: True

oneflow/core/functional/impl/math_functor.cpp

@@ -447,6 +447,67 @@ class MaximumFunctor {
   std::shared_ptr<OpExpr> broadcast_maximum_op_;
 };
 
+class ScalarLogicalBaseFunctor {
+ public:
+  explicit ScalarLogicalBaseFunctor(std::string op_name) {
+    op_ = CHECK_JUST(one::OpBuilder(op_name).Input("in").Output("out").Build());
+  }
+  virtual ~ScalarLogicalBaseFunctor() = default;
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const Scalar& scalar) const {
+    MutableAttrMap attrs;
+
+    if (IsFloatingDataType(x->dtype())) {
+      JUST(attrs.SetAttr<double>("float_operand", JUST(scalar.As<double>())));
+      JUST(attrs.SetAttr<bool>("has_float_operand", true));
+      JUST(attrs.SetAttr<bool>("has_int_operand", false));
+    } else if (IsIntegralDataType(x->dtype())) {
+      JUST(attrs.SetAttr<int64_t>("int_operand", JUST(scalar.As<int64_t>())));
+      JUST(attrs.SetAttr<bool>("has_float_operand", false));
+      JUST(attrs.SetAttr<bool>("has_int_operand", true));
+    } else {
+      UNIMPLEMENTED_THEN_RETURN() << "The scalar in ScalarAdd shoule be float or int.";
+    }
+
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {x}, attrs);
+  }
+
+ private:
+  std::shared_ptr<OpExpr> op_;
+};
+
+class ScalarLogicalEqualFunctor : public ScalarLogicalBaseFunctor {
+ public:
+  ScalarLogicalEqualFunctor() : ScalarLogicalBaseFunctor(/*op_name=*/"scalar_logical_equal") {}
+};
+
+class ScalarLogicalNotEqualFunctor : public ScalarLogicalBaseFunctor {
+ public:
+  ScalarLogicalNotEqualFunctor()
+      : ScalarLogicalBaseFunctor(/*op_name=*/"scalar_logical_not_equal") {}
+};
+
+class ScalarLogicalGreaterFunctor : public ScalarLogicalBaseFunctor {
+ public:
+  ScalarLogicalGreaterFunctor() : ScalarLogicalBaseFunctor(/*op_name=*/"scalar_logical_greater") {}
+};
+
+class ScalarLogicalGreaterEqualFunctor : public ScalarLogicalBaseFunctor {
+ public:
+  ScalarLogicalGreaterEqualFunctor()
+      : ScalarLogicalBaseFunctor(/*op_name=*/"scalar_logical_greater_equal") {}
+};
+
+class ScalarLogicalLessFunctor : public ScalarLogicalBaseFunctor {
+ public:
+  ScalarLogicalLessFunctor() : ScalarLogicalBaseFunctor(/*op_name=*/"scalar_logical_less") {}
+};
+
+class ScalarLogicalLessEqualFunctor : public ScalarLogicalBaseFunctor {
+ public:
+  ScalarLogicalLessEqualFunctor()
+      : ScalarLogicalBaseFunctor(/*op_name=*/"scalar_logical_less_equal") {}
+};
+
 }  // namespace impl
 
 ONEFLOW_FUNCTION_LIBRARY(m) {
@@ -469,6 +530,12 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::SelectFirstFunctor>("SelectFirst");
   m.add_functor<impl::MinimumFunctor>("Minimum");
   m.add_functor<impl::MaximumFunctor>("Maximum");
+  m.add_functor<impl::ScalarLogicalEqualFunctor>("ScalarLogicalEqual");
+  m.add_functor<impl::ScalarLogicalNotEqualFunctor>("ScalarLogicalNotEqual");
+  m.add_functor<impl::ScalarLogicalGreaterFunctor>("ScalarLogicalGreater");
+  m.add_functor<impl::ScalarLogicalGreaterEqualFunctor>("ScalarLogicalGreaterEqual");
+  m.add_functor<impl::ScalarLogicalLessFunctor>("ScalarLogicalLess");
+  m.add_functor<impl::ScalarLogicalLessEqualFunctor>("ScalarLogicalLessEqual");
 };
 
 }  // namespace functional