Finish T1-1-11: gcd、select_scatter、nll_loss、glu、gt

Author: PPPoint-t

include/infinicore/ops/gcd.hpp

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

include/infinicore/ops/glu.hpp

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

include/infinicore/ops/gt.hpp

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

include/infinicore/ops/nll_loss.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+class NLLLoss {
+public:
+    using schema = void (*)(Tensor, Tensor, std::optional<Tensor>, Tensor, int64_t);
+    static void execute(Tensor input, Tensor target, std::optional<Tensor> weight, Tensor output, int64_t ignore_index);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor nll_loss(Tensor input, Tensor target, std::optional<Tensor> weight, int64_t ignore_index);
+void nll_loss_(Tensor input, Tensor target, std::optional<Tensor> weight, Tensor output, int64_t ignore_index);
+
+} // namespace infinicore::op

include/infinicore/ops/select_scatter.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class SelectScatter {
+public:
+    using schema = void (*)(Tensor, Tensor, int64_t, int64_t, Tensor);
+    static void execute(Tensor input, Tensor src, int64_t dim, int64_t index, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor select_scatter(Tensor input, Tensor src, int64_t dim, int64_t index);
+void select_scatter_(Tensor input, Tensor src, int64_t dim, int64_t index, Tensor output);
+
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -54,6 +54,9 @@
 from infinicore.ops.rearrange import rearrange
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.unsqueeze import unsqueeze
+from infinicore.ops.gcd import gcd
+from infinicore.ops.gt import gt
+from infinicore.ops.select_scatter import select_scatter
 from infinicore.tensor import (
     Tensor,
     empty,
@@ -134,6 +137,9 @@
     "strided_empty",
     "strided_from_blob",
     "zeros",
+    "gcd",
+    "select_scatter",
+    "gt",
 ]
 
 use_ntops = False

python/infinicore/nn/functional/__init__.py

@@ -6,6 +6,8 @@
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
+from .nll_loss import nll_loss
+from .glu import glu
 
 __all__ = [
     "causal_softmax",
@@ -17,4 +19,6 @@
     "embedding",
     "rope",
     "RopeAlgo",
+    "nll_loss",
+    "glu",
 ]

python/infinicore/nn/functional/glu.py

@@ -0,0 +1,10 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def glu(input: Tensor, dim: int = -1) -> Tensor:
+    
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.glu(input, dim)
+
+    return Tensor(_infinicore.glu(input._underlying, dim))

python/infinicore/nn/functional/nll_loss.py

@@ -0,0 +1,39 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def nll_loss(
+    input: Tensor,
+    target: Tensor,
+    weight: Tensor | None = None,
+    ignore_index: int = -100,
+    reduction: str = "mean",
+    *,
+    out=None,
+) -> Tensor:
+    
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.nll_loss(
+            input, target, weight=weight, ignore_index=ignore_index, reduction=reduction
+        )
+
+    weight_underlying = weight._underlying if weight is not None else None
+
+    if out is None:
+        return Tensor(
+            _infinicore.nll_loss(
+                input._underlying, 
+                target._underlying, 
+                weight_underlying, 
+                ignore_index
+            )
+        )
+    
+    _infinicore.nll_loss_(
+        input._underlying, 
+        target._underlying, 
+        weight_underlying, 
+        out._underlying, 
+        ignore_index
+    )
+    return out

python/infinicore/ops/gcd.py

@@ -0,0 +1,15 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def gcd(input: Tensor, other: Tensor, *, out=None) -> Tensor:
+    r"""Computes the element-wise greatest common divisor (GCD)."""
+    
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.gcd(input, other, out=out)
+
+    if out is None:
+        return Tensor(_infinicore.gcd(input._underlying, other._underlying))
+    
+    _infinicore.gcd_(input._underlying, other._underlying, out._underlying)
+    return out