[2025秋季][T1-1-32] JBFYS-XD

include/infinicore/ops/cat.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+Tensor cat(std::vector<Tensor> tensors, int dim);
+void cat_(Tensor out, std::vector<Tensor> tensors, int dim);
+} // namespace infinicore::op

include/infinicore/ops/inner.hpp

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

include/infinicore/ops/masked_select.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class MaskedSelect {
+public:
+    using schema = void (*)(Tensor, Tensor, void **, size_t *);
+    static void execute(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor masked_select(Tensor input, Tensor mask);
+
+} // namespace infinicore::op

include/infinicore/ops/tan.hpp

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

include/infinicore/ops/tanhshrink.hpp

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

include/infiniop.h

@@ -10,9 +10,11 @@
 #include "infiniop/ops/dequantize_awq.h"
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gemm.h"
+#include "infiniop/ops/inner.h"
 #include "infiniop/ops/layer_norm.h"
 #include "infiniop/ops/logsoftmax.h"
 #include "infiniop/ops/lp_norm.h"
+#include "infiniop/ops/masked_select.h"
 #include "infiniop/ops/mul.h"
 #include "infiniop/ops/ones.h"
 #include "infiniop/ops/random_sample.h"
@@ -26,7 +28,9 @@
 #include "infiniop/ops/softplus.h"
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
+#include "infiniop/ops/tan.h"
 #include "infiniop/ops/tanh.h"
+#include "infiniop/ops/tanhshrink.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"
 #include "infiniop/ops/zeros.h"

include/infiniop/ops/inner.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_INNER_API_H__
+#define __INFINIOP_INNER_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopInnerDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateInnerDescriptor(
+    infiniopHandle_t handle,
+    infiniopInnerDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t other_desc);
+
+__C __export infiniStatus_t infiniopGetInnerWorkspaceSize(infiniopInnerDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopInner(
+    infiniopInnerDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *input,
+    const void *other,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyInnerDescriptor(infiniopInnerDescriptor_t desc);
+
+#endif

include/infiniop/ops/masked_select.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_MASKED_SELECT_API_H__
+#define __INFINIOP_MASKED_SELECT_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopMaskedSelectDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateMaskedSelectDescriptor(
+    infiniopHandle_t handle,
+    infiniopMaskedSelectDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t mask_desc);
+
+__C __export infiniStatus_t infiniopGetMaskedSelectWorkspaceSize(infiniopMaskedSelectDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopMaskedSelect(
+    infiniopMaskedSelectDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    const void *input,
+    const bool *mask,
+    void **data_ptr,
+    size_t *dlen,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyMaskedSelectDescriptor(infiniopMaskedSelectDescriptor_t desc);
+
+#endif

include/infiniop/ops/tan.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_TAN_API_H__
+#define __INFINIOP_TAN_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTanDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTanDescriptor(
+    infiniopHandle_t handle,
+    infiniopTanDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t out_desc);
+
+__C __export infiniStatus_t infiniopGetTanWorkspaceSize(infiniopTanDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTan(
+    infiniopTanDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    const void *input,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyTanDescriptor(infiniopTanDescriptor_t desc);
+
+#endif

include/infiniop/ops/tanhshrink.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_TANHSHRINK_API_H__
+#define __INFINIOP_TANHSHRINK_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTanhshrinkDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTanhshrinkDescriptor(infiniopHandle_t handle,
+                                                         infiniopTanhshrinkDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t output,
+                                                         infiniopTensorDescriptor_t intput);
+
+__C __export infiniStatus_t infiniopGetTanhshrinkWorkspaceSize(infiniopTanhshrinkDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTanhshrink(infiniopTanhshrinkDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *output,
+                                         const void *intput,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyTanhshrinkDescriptor(infiniopTanhshrinkDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -41,10 +41,14 @@
 )
 from infinicore.ops.add import add
 from infinicore.ops.attention import attention
+from infinicore.ops.cat import cat
+from infinicore.ops.inner import inner
+from infinicore.ops.masked_select import masked_select
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
 from infinicore.ops.rearrange import rearrange
+from infinicore.ops.tan import tan
 from infinicore.tensor import (
     Tensor,
     empty,
@@ -101,10 +105,14 @@
     # Operations.
     "add",
     "attention",
+    "cat",
+    "inner",
+    "masked_select",
     "matmul",
     "mul",
     "narrow",
     "rearrange",
+    "tan",
     "empty",
     "empty_like",
     "from_blob",

python/infinicore/nn/functional/__init__.py

@@ -6,6 +6,7 @@
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
+from .tanhshrink import tanhshrink
 
 __all__ = [
     "causal_softmax",
@@ -17,4 +18,5 @@
     "embedding",
     "rope",
     "RopeAlgo",
+    "tanhshrink"
 ]

python/infinicore/nn/functional/tanhshrink.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def tanhshrink(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.tanhshrink(input._underlying))
+
+    _infinicore.tanhshrink_(out._underlying, input._underlying)
+
+    return out

python/infinicore/ops/cat.py

@@ -0,0 +1,12 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def cat(tensors, dim:int=0, *,out=None):
+    if out is None:
+        return Tensor(_infinicore.cat([tensor._underlying for tensor in tensors], dim))
+
+    _infinicore.cat_(out._underlying, [tensor._underlying for tensor in tensors], dim)
+
+    # raise RuntimeError("breakpointer!!!")
+    return out

python/infinicore/ops/inner.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def inner(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.inner(input._underlying, other._underlying))
+
+    _infinicore.inner_(out._underlying, input._underlying, other._underlying)
+
+    return out

python/infinicore/ops/masked_select.py

@@ -0,0 +1,5 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def masked_select(input, mask):
+    return Tensor(_infinicore.masked_select(input._underlying, mask._underlying))

python/infinicore/ops/tan.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def tan(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.tan(input._underlying))
+
+    _infinicore.tan_(out._underlying, input._underlying)
+
+    return out