Revert "Merge pull request #1056 from InfiniTensor/issue/1031"

This reverts commit 7f295448, reversing
changes made to e60985dc.

include/infinicore/ops.hpp

@@ -3,13 +3,9 @@
 #include "ops/add.hpp"
 #include "ops/add_rms_norm.hpp"
 #include "ops/attention.hpp"
-#include "ops/avg_pool1d.hpp"
 #include "ops/causal_softmax.hpp"
-#include "ops/cross_entropy.hpp"
 #include "ops/embedding.hpp"
 #include "ops/flash_attention.hpp"
-#include "ops/hardswish.hpp"
-#include "ops/hardtanh.hpp"
 #include "ops/kv_caching.hpp"
 #include "ops/matmul.hpp"
 #include "ops/ones.hpp"

include/infinicore/ops/avg_pool1d.hpp

-#pragma once
-
-#include "../device.hpp"
-#include "common/op.hpp"
-
-namespace infinicore::op {
-
-class AvgPool1d {
-public:
-    using schema = void (*)(Tensor, Tensor, size_t, size_t, size_t);
-    static void execute(Tensor output, Tensor input, size_t kernel_size, size_t stride, size_t padding);
-    static common::OpDispatcher<schema> &dispatcher();
-};
-
-Tensor avg_pool1d(Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0);
-void avg_pool1d_(Tensor output, Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0);
-
-} // namespace infinicore::op

include/infinicore/ops/cross_entropy.hpp

-#pragma once
-
-#include "../device.hpp"
-#include "common/op.hpp"
-
-namespace infinicore::op {
-
-class CrossEntropy {
-public:
-    // Schema 定义：函数指针类型
-    // CrossEntropy 需要接收三个 Tensor: Output (Loss), Input (Logits), Target (Labels)
-    using schema = void (*)(Tensor, Tensor, Tensor);
-
-    // 执行入口
-    static void execute(Tensor output, Tensor input, Tensor target);
-
-    // 分发器访问接口
-    static common::OpDispatcher<schema> &dispatcher();
-};
-
-// ==================================================================
-// 对外 Functional API
-// ==================================================================
-
-// 1. Out-of-place 接口：
-// 输入 Logits 和 Target，内部自动创建 Output Tensor 并返回
-Tensor cross_entropy(Tensor input, Tensor target);
-
-// 2. Explicit Output 接口 (类似于 In-place 风格)：
-// 用户显式提供 Output Tensor 用于存储结果
-// 注意：虽然命名带有下划线 _，但通常 CrossEntropy 无法真正原地修改 input，
-// 所以这里只是表示“写入指定的 output 内存”
-void cross_entropy_(Tensor output, Tensor input, Tensor target);
-
-} // namespace infinicore::op

include/infinicore/ops/equal.hpp

-#pragma once
-
-#include "../device.hpp"
-#include "common/op.hpp"
-
-namespace infinicore::op {
-
-class Equal {
-public:
-    using schema = void (*)(Tensor, Tensor, Tensor);
-
-    static void execute(Tensor out, Tensor a, Tensor b);
-    static common::OpDispatcher<schema> &dispatcher();
-};
-
-Tensor equal(Tensor a, Tensor b);
-void equal_(Tensor out, Tensor a, Tensor b);
-
-} // namespace infinicore::op

include/infinicore/ops/hardswish.hpp

-#pragma once
-
-#include "../device.hpp"
-#include "common/op.hpp"
-
-namespace infinicore::op {
-
-class Hardswish {
-public:
-    using schema = void (*)(Tensor, Tensor);
-    static void execute(Tensor output, Tensor input);
-    static common::OpDispatcher<schema> &dispatcher();
-};
-
-Tensor hardswish(Tensor input);
-void hardswish_(Tensor output, Tensor input);
-
-} // namespace infinicore::op

include/infinicore/ops/hardtanh.hpp

-#pragma once
-
-#include "../device.hpp"
-#include "common/op.hpp"
-
-namespace infinicore::op {
-
-class HardTanh {
-public:
-    using schema = void (*)(Tensor, Tensor, float, float);
-    static void execute(Tensor output, Tensor input, float min_val, float max_val);
-    static common::OpDispatcher<schema> &dispatcher();
-};
-
-Tensor hardtanh(Tensor input, float min_val = -1.0f, float max_val = 1.0f);
-void hardtanh_(Tensor output, Tensor input, float min_val = -1.0f, float max_val = 1.0f);
-
-} // namespace infinicore::op

include/infiniop.h

@@ -42,10 +42,4 @@
 #include "infiniop/ops/zeros.h"
 #include "infiniop/tensor_descriptor.h"
 
-#include "infiniop/ops/cross_entropy.h"
-#include "infiniop/ops/hardswish.h"
-#include "infiniop/ops/avg_pool1d.h"
-#include "infiniop/ops/equal.h"
-#include "infiniop/ops/hardtanh.h"
-
 #endif // __INFINIOP_API_H__

include/infiniop/ops/avg_pool1d.h

-#ifndef __INFINIOP_AVG_POOL1D_API_H__
-#define __INFINIOP_AVG_POOL1D_API_H__
-
-#include "../operator_descriptor.h"
-
-typedef struct InfiniopDescriptor *infiniopAvgPool1dDescriptor_t;
-
-__INFINI_C __export infiniStatus_t infiniopCreateAvgPool1dDescriptor(
-    infiniopHandle_t handle,
-    infiniopAvgPool1dDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t output,
-    infiniopTensorDescriptor_t input,
-    size_t kernel_size,
-    size_t stride,
-    size_t padding);
-
-__INFINI_C __export infiniStatus_t infiniopGetAvgPool1dWorkspaceSize(
-    infiniopAvgPool1dDescriptor_t desc,
-    size_t *size);
-
-__INFINI_C __export infiniStatus_t infiniopAvgPool1d(
-    infiniopAvgPool1dDescriptor_t desc,
-    void *workspace,
-    size_t workspace_size,
-    void *output,
-    const void *input,
-    void *stream);
-
-__INFINI_C __export infiniStatus_t infiniopDestroyAvgPool1dDescriptor(
-    infiniopAvgPool1dDescriptor_t desc);
-
-#endif

include/infiniop/ops/cross_entropy.h

-#ifndef __INFINIOP_CROSS_ENTROPY_API_H__
-#define __INFINIOP_CROSS_ENTROPY_API_H__
-
-#include "../operator_descriptor.h"
-
-typedef struct InfiniopDescriptor *infiniopCrossEntropyDescriptor_t;
-
-__INFINI_C __export infiniStatus_t infiniopCreateCrossEntropyDescriptor(
-    infiniopHandle_t handle,
-    infiniopCrossEntropyDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t y_desc,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t target_desc);
-
-__INFINI_C __export infiniStatus_t infiniopGetCrossEntropyWorkspaceSize(
-    infiniopCrossEntropyDescriptor_t desc,
-    size_t *size);
-
-__INFINI_C __export infiniStatus_t infiniopCrossEntropy(
-    infiniopCrossEntropyDescriptor_t desc,
-    void *workspace,
-    size_t workspace_size,
-    void *y,
-    const void *x,
-    const void *target,
-    void *stream);
-
-__INFINI_C __export infiniStatus_t infiniopDestroyCrossEntropyDescriptor(
-    infiniopCrossEntropyDescriptor_t desc);
-
-#endif

include/infiniop/ops/equal.h

-#ifndef __INFINIOP_EQUAL_API_H__
-#define __INFINIOP_EQUAL_API_H__
-
-#include "../operator_descriptor.h"
-
-typedef struct InfiniopDescriptor *infiniopEqualDescriptor_t;
-
-__INFINI_C __export infiniStatus_t infiniopCreateEqualDescriptor(
-    infiniopHandle_t handle,
-    infiniopEqualDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t c,
-    infiniopTensorDescriptor_t a,
-    infiniopTensorDescriptor_t b);
-
-__INFINI_C __export infiniStatus_t infiniopGetEqualWorkspaceSize(
-    infiniopEqualDescriptor_t desc,
-    size_t *size);
-
-__INFINI_C __export infiniStatus_t infiniopEqual(
-    infiniopEqualDescriptor_t desc,
-    void *workspace,
-    size_t workspace_size,
-    void *c,
-    const void *a,
-    const void *b,
-    void *stream);
-
-__INFINI_C __export infiniStatus_t infiniopDestroyEqualDescriptor(
-    infiniopEqualDescriptor_t desc);
-
-#endif

include/infiniop/ops/hardswish.h

-#ifndef __INFINIOP_HARDSWISH_API_H__
-#define __INFINIOP_HARDSWISH_API_H__
-
-#include "../operator_descriptor.h"
-
-typedef struct InfiniopDescriptor *infiniopHardSwishDescriptor_t;
-
-__INFINI_C __export infiniStatus_t infiniopCreateHardSwishDescriptor(
-    infiniopHandle_t handle,
-    infiniopHardSwishDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t output,
-    infiniopTensorDescriptor_t input);
-
-__INFINI_C __export infiniStatus_t infiniopGetHardSwishWorkspaceSize(
-    infiniopHardSwishDescriptor_t desc,
-    size_t *size);
-
-__INFINI_C __export infiniStatus_t infiniopHardSwish(
-    infiniopHardSwishDescriptor_t desc,
-    void *workspace,
-    size_t workspace_size,
-    void *output,
-    const void *input,
-    void *stream);
-
-__INFINI_C __export infiniStatus_t infiniopDestroyHardSwishDescriptor(
-    infiniopHardSwishDescriptor_t desc);
-
-#endif

include/infiniop/ops/hardtanh.h

-#ifndef __INFINIOP_HARDTANH_API_H__
-#define __INFINIOP_HARDTANH_API_H__
-
-#include "../operator_descriptor.h"
-
-typedef struct InfiniopDescriptor *infiniopHardTanhDescriptor_t;
-
-__INFINI_C __export infiniStatus_t infiniopCreateHardTanhDescriptor(infiniopHandle_t handle,
-                                                                    infiniopHardTanhDescriptor_t *desc_ptr,
-                                                                    infiniopTensorDescriptor_t output,
-                                                                    infiniopTensorDescriptor_t input,
-                                                                    float min_val,
-                                                                    float max_val);
-
-__INFINI_C __export infiniStatus_t infiniopGetHardTanhWorkspaceSize(infiniopHardTanhDescriptor_t desc,
-                                                                    size_t *size);
-
-__INFINI_C __export infiniStatus_t infiniopHardTanh(infiniopHardTanhDescriptor_t desc,
-                                                    void *workspace,
-                                                    size_t workspace_size,
-                                                    void *output,
-                                                    const void *input,
-                                                    void *stream);
-
-__INFINI_C __export infiniStatus_t infiniopDestroyHardTanhDescriptor(infiniopHardTanhDescriptor_t desc);
-
-#endif

python/infinicore/__init__.py

@@ -50,8 +50,6 @@ from infinicore.dtype import (
 from infinicore.ops.add import add
 from infinicore.ops.add_rms_norm import add_rms_norm
 from infinicore.ops.attention import attention
-from infinicore.ops.cross_entropy import cross_entropy
-from infinicore.ops.equal import equal
 from infinicore.ops.kv_caching import kv_caching
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mha_kvcache import mha_kvcache
@@ -127,13 +125,11 @@ __all__ = [
     "attention",
     "kv_caching",
     "matmul",
-    "equal",
     "mul",
     "narrow",
     "squeeze",
     "unsqueeze",
     "rearrange",
-    "cross_entropy",
     "empty",
     "empty_like",
     "from_blob",

python/infinicore/nn/functional/__init__.py

-from .avg_pool1d import avg_pool1d
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .flash_attention import flash_attention
-from .hardswish import hardswish
-from .hardtanh import hardtanh
 from .linear import linear
 from .linear_w8a8i8 import linear_w8a8i8
 from .random_sample import random_sample
@@ -23,9 +20,6 @@ __all__ = [
     "RopeAlgo",
     "rope",
     "silu",
-    "hardswish",
-    "hardtanh",
-    "avg_pool1d",
     "swiglu",
     "linear_w8a8i8",
     "silu_and_mul",

python/infinicore/nn/functional/avg_pool1d.py

-from infinicore.lib import _infinicore
-from infinicore.tensor import Tensor
-
-
-def avg_pool1d(
-    input: Tensor,
-    kernel_size: int,
-    stride: int | None = None,
-    padding: int = 0,
-    *,
-    out=None,
-) -> Tensor:
-    if stride is None:
-        stride = 0
-
-    if out is None:
-        return Tensor(
-            _infinicore.avg_pool1d(input._underlying, kernel_size, stride, padding)
-        )
-
-    _infinicore.avg_pool1d_(
-        out._underlying, input._underlying, kernel_size, stride, padding
-    )
-    return out

python/infinicore/nn/functional/hardswish.py

-import infinicore
-from infinicore.lib import _infinicore
-from infinicore.tensor import Tensor
-
-
-def hardswish(input: Tensor, inplace: bool = False, *, out=None) -> Tensor:
-    r"""Apply the Hardswish activation function element-wise."""
-
-    if (
-        infinicore.use_ntops
-        and input.device.type in ("cuda", "musa")
-        and out is None
-        and hasattr(infinicore.ntops.torch, "hardswish")
-    ):
-        try:
-            return infinicore.ntops.torch.hardswish(input, inplace=inplace)
-        except AttributeError:
-            pass
-
-    if inplace:
-        _infinicore.hardswish_(input._underlying, input._underlying)
-        return input
-
-    if out is None:
-        return Tensor(_infinicore.hardswish(input._underlying))
-
-    _infinicore.hardswish_(out._underlying, input._underlying)
-    return out

python/infinicore/nn/functional/hardtanh.py

-import infinicore
-from infinicore.lib import _infinicore
-from infinicore.tensor import Tensor
-
-
-def hardtanh(
-    input: Tensor,
-    min_val: float = -1.0,
-    max_val: float = 1.0,
-    inplace: bool = False,
-    *,
-    out=None,
-) -> Tensor:
-    """Clamp the input tensor to the range [min_val, max_val]."""
-
-    if min_val > max_val:
-        raise ValueError("min_val must be less than or equal to max_val")
-
-    if (
-        infinicore.use_ntops
-        and input.device.type in ("cuda", "musa")
-        and out is None
-        and hasattr(infinicore.ntops.torch, "hardtanh")
-    ):
-        try:
-            return infinicore.ntops.torch.hardtanh(
-                input, min_val=min_val, max_val=max_val, inplace=inplace
-            )
-        except AttributeError:
-            pass
-
-    if inplace:
-        _infinicore.hardtanh_(
-            input._underlying, input._underlying, float(min_val), float(max_val)
-        )
-        return input
-
-    if out is None:
-        return Tensor(
-            _infinicore.hardtanh(input._underlying, float(min_val), float(max_val))
-        )
-
-    _infinicore.hardtanh_(
-        out._underlying, input._underlying, float(min_val), float(max_val)
-    )
-    return out

python/infinicore/ops/cross_entropy.py

-from infinicore.lib import _infinicore
-from infinicore.tensor import Tensor
-
-
-def cross_entropy(
-    logits,
-    target,
-    weight=None,
-    *,
-    ignore_index=None,
-    reduction="none",
-    out=None,
-):
-    """
-    Token-wise cross entropy without reduction. The output tensor has the same
-    shape as target and uses the logits dtype.
-    """
-    if weight is not None:
-        raise NotImplementedError("class weights are not supported yet.")
-    if ignore_index is not None:
-        raise NotImplementedError("ignore_index is not supported yet.")
-    if reduction not in (None, "none"):
-        raise NotImplementedError("Only reduction='none' is implemented.")
-
-    if out is None:
-        return Tensor(_infinicore.cross_entropy(logits._underlying, target._underlying))
-
-    _infinicore.cross_entropy_(
-        out._underlying,
-        logits._underlying,
-        target._underlying,
-    )
-    return out

python/infinicore/ops/equal.py

-from infinicore.lib import _infinicore
-from infinicore.tensor import Tensor
-
-
-def equal(input, other, *, out=None):
-    if out is None:
-        return Tensor(_infinicore.equal(input._underlying, other._underlying))
-
-    _infinicore.equal_(out._underlying, input._underlying, other._underlying)
-    return out

python/infinicore/utils.py

+import ml_dtypes
 import numpy as np
 import torch
 
 import infinicore
 
-try:
-    import ml_dtypes
-except ModuleNotFoundError:
-    ml_dtypes = None
-
 
 def to_torch_dtype(infini_dtype):
     """Convert infinicore data type to PyTorch data type"""
@@ -61,9 +57,7 @@ def numpy_to_infinicore_dtype(numpy_dtype):
         return infinicore.float64
     elif numpy_dtype == np.float16:
         return infinicore.float16
-    elif hasattr(np, "bfloat16") and numpy_dtype == np.bfloat16:
-        return infinicore.bfloat16
-    elif ml_dtypes is not None and numpy_dtype == ml_dtypes.bfloat16:
+    elif numpy_dtype == ml_dtypes.bfloat16:
         return infinicore.bfloat16
     elif numpy_dtype == np.int8:
         return infinicore.int8
@@ -92,13 +86,6 @@ def infinicore_to_numpy_dtype(infini_dtype):
     elif infini_dtype == infinicore.int16:
         return np.int16
     elif infini_dtype == infinicore.bfloat16:
-        if hasattr(np, "bfloat16"):
-            return np.bfloat16
-        if ml_dtypes is None:
-            raise ModuleNotFoundError(
-                "ml_dtypes is required for bfloat16 numpy conversion. "
-                "Please install ml_dtypes."
-            )
         return ml_dtypes.bfloat16
     elif infini_dtype == infinicore.int32:
         return np.int32