Finish T1-1-12: minimum、atan2、addcdiv、bucketize、binary_cross_entropy

Author: PPPoint-t

include/infinicore/ops/addcdiv.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Addcdiv {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float);
+    static void execute(Tensor input, Tensor t1, Tensor t2, Tensor output, float value);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor addcdiv(Tensor input, Tensor t1, Tensor t2, float value);
+void addcdiv_(Tensor input, Tensor t1, Tensor t2, Tensor output, float value);
+
+} // namespace infinicore::op

include/infinicore/ops/atan2.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Atan2 {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor input, Tensor other, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor atan2(Tensor input, Tensor other);
+void atan2_(Tensor input, Tensor other, Tensor output);
+
+} // namespace infinicore::op

include/infinicore/ops/binary_cross_entropy.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+#include <string>
+
+namespace infinicore::op {
+
+class BinaryCrossEntropy {
+public:
+    using schema = void (*)(Tensor, Tensor, std::optional<Tensor>, Tensor, std::string);
+    static void execute(Tensor input, Tensor target, std::optional<Tensor> weight, Tensor output, std::string reduction);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor binary_cross_entropy(Tensor input, Tensor target, std::optional<Tensor> weight, std::string reduction);
+void binary_cross_entropy_(Tensor input, Tensor target, std::optional<Tensor> weight, Tensor output, std::string reduction);
+
+} // namespace infinicore::op

include/infinicore/ops/bucketize.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Bucketize {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, bool);
+    static void execute(Tensor input, Tensor boundaries, Tensor output, bool right);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor bucketize(Tensor input, Tensor boundaries, bool right = false);
+void bucketize_(Tensor input, Tensor boundaries, Tensor output, bool right = false);
+
+} // namespace infinicore::op

include/infinicore/ops/minimum.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Minimum {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor input, Tensor other, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor minimum(Tensor input, Tensor other);
+void minimum_(Tensor input, Tensor other, Tensor output);
+
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -44,8 +44,12 @@
 )
 from infinicore.ops.add import add
 from infinicore.ops.add_rms_norm import add_rms_norm, add_rms_norm_
+from infinicore.ops.addcdiv import addcdiv
+from infinicore.ops.atan2 import atan2
 from infinicore.ops.attention import attention
+from infinicore.ops.bucketize import bucketize
 from infinicore.ops.matmul import matmul
+from infinicore.ops.minimum import minimum
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
 from infinicore.ops.paged_attention import paged_attention

python/infinicore/nn/functional/__init__.py

@@ -1,3 +1,4 @@
+from .binary_cross_entropy import binary_cross_entropy
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .linear import linear
@@ -17,4 +18,5 @@
     "embedding",
     "rope",
     "RopeAlgo",
+    "binary_cross_entropy",
 ]

python/infinicore/nn/functional/binary_cross_entropy.py

@@ -0,0 +1,47 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def binary_cross_entropy(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    size_average=None,
+    reduce=None,
+    reduction: str = "mean",
+    *,
+    out=None,
+) -> Tensor:
+    r"""Apply the binary_cross_entropy function."""
+
+    if size_average is not None or reduce is not None:
+        if reduce is False:
+            reduction = "none"
+        elif size_average is True or size_average is None:
+            reduction = "mean"
+        else:
+            reduction = "sum"
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.binary_cross_entropy(
+            input, target, weight=weight, reduction=reduction, out=out
+        )
+
+    weight_underlying = weight._underlying if weight is not None else None
+
+    if out is None:
+        return Tensor(
+            _infinicore.binary_cross_entropy(
+                input._underlying, target._underlying, weight_underlying, reduction
+            )
+        )
+
+    _infinicore.binary_cross_entropy_(
+        input._underlying,
+        target._underlying,
+        weight_underlying,
+        out._underlying,
+        reduction,
+    )
+    return out

python/infinicore/ops/addcdiv.py

@@ -0,0 +1,38 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def addcdiv(
+    input: Tensor,
+    tensor1: Tensor,
+    tensor2: Tensor,
+    *,
+    value=1.0,
+    out=None,
+) -> Tensor:
+    r"""Apply the addcdiv function."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.addcdiv(
+            input, tensor1, tensor2, value=value, out=out
+        )
+
+    if out is None:
+        return Tensor(
+            _infinicore.addcdiv(
+                input._underlying,
+                tensor1._underlying,
+                tensor2._underlying,
+                float(value),
+            )
+        )
+
+    _infinicore.addcdiv_(
+        input._underlying,
+        tensor1._underlying,
+        tensor2._underlying,
+        out._underlying,
+        float(value),
+    )
+    return out

python/infinicore/ops/atan2.py

@@ -0,0 +1,21 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def atan2(
+    input: Tensor,
+    other: Tensor,
+    *,
+    out=None,
+) -> Tensor:
+    r"""Apply the atan2 function."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.atan2(input, other, out=out)
+
+    if out is None:
+        return Tensor(_infinicore.atan2(input._underlying, other._underlying))
+
+    _infinicore.atan2_(input._underlying, other._underlying, out._underlying)
+    return out