issue/892-添加infini.tensor函数; 移除from_numpy函数的dtype参数

Author: pengcheng888

python/infinicore/__init__.py

@@ -62,6 +62,7 @@
     ones,
     strided_empty,
     strided_from_blob,
+    tensor,
     zeros,
 )
 
@@ -127,6 +128,7 @@
     "ones",
     "strided_empty",
     "strided_from_blob",
+    "tensor",
     "zeros",
 ]
 

python/infinicore/nn/modules/rope.py

@@ -4,6 +4,7 @@
 from infinicore.nn import functional as F
 
 from ...tensor import Tensor
+from ...utils import infinicore_to_numpy_dtype
 from ..functional import RopeAlgo
 from .module import InfiniCoreModule as Module
 
@@ -31,8 +32,13 @@ def create_sin_cos_table(
         max_position, head_dim, theta
     )
 
-    sin_table_infini = infinicore.from_numpy(sin_table_np, dtype=dtype, device=device)
-    cos_table_infini = infinicore.from_numpy(cos_table_np, dtype=dtype, device=device)
+    if dtype is not None:
+        np_dtype = infinicore_to_numpy_dtype(dtype)
+        sin_table_np = sin_table_np.astype(np_dtype)
+        cos_table_np = cos_table_np.astype(np_dtype)
+
+    sin_table_infini = infinicore.from_numpy(sin_table_np, device=device)
+    cos_table_infini = infinicore.from_numpy(cos_table_np, device=device)
 
     return sin_table_infini, cos_table_infini
 

python/infinicore/tensor.py

@@ -1,4 +1,5 @@
 import ctypes
+from typing import Any, Union
 
 import numpy as np
 
@@ -17,15 +18,18 @@ class Tensor:
     # Public attributes describing the Tensor
     _underlying: _infinicore.Tensor
     _torch_ref: "torch.Tensor"  # noqa: F821
+    _numpy_ref: np.ndarray  # noqa: F821
+
     shape: list[int]
     dtype: infinicore.dtype
     device: infinicore.device
 
-    def __init__(self, underlying, *, _torch_ref=None):
+    def __init__(self, underlying, *, _torch_ref=None, _numpy_ref=None):
         """An internal method. Please do not use this directly."""
 
         self._underlying = underlying
         self._torch_ref = _torch_ref
+        self._numpy_ref = _numpy_ref
 
     def __getattr__(self, name):
         # Lazily construct and cache an attribute.
@@ -131,6 +135,38 @@ def narrow(self, dim, start, length):
         return infinicore.narrow(self, dim, start, length)
 
 
+def tensor(
+    data: Any,
+    *,
+    dtype: Union[infinicore.dtype, None] = None,
+    device: Union[infinicore.device, str, int, None] = None,
+    pin_memory: bool = False,
+) -> Tensor:
+    r"""
+    Constructs a tensor by copying `data`.
+
+    Args:
+        data (array_like): Initial data for the tensor. Can be a list, tuple, NumPy, scalar.
+
+    Keyword args:
+        dtype (infinicore.dtype, optional): the desired data type of returned tensor.
+            Default: if ``None``, infers data type from :attr:`data`.
+        device (infinicore.device, optional): the device of the constructed tensor.
+    """
+
+    if isinstance(data, (list, tuple)):
+        return from_list(data, dtype=dtype, device=device)
+    elif isinstance(data, np.ndarray):
+        if dtype is not None:
+            np_dtype = infinicore_to_numpy_dtype(dtype)
+            data = data.astype(np_dtype)
+        return from_numpy(data, device=device)
+    elif isinstance(data, (int, float)):
+        return from_list([data], dtype=dtype, device=device).squeeze(0)
+
+    raise ValueError(f" Can not convert data type: {type(data)} to Tensor")
+
+
 def empty(size, *, dtype=None, device=None, pin_memory=False):
     return Tensor(
         _infinicore.empty(size, dtype._underlying, device._underlying, pin_memory)
@@ -198,28 +234,27 @@ def from_torch(torch_tensor) -> Tensor:
 def from_numpy(
     np_array,
     *,
-    dtype: infinicore.dtype = None,
     device: infinicore.device = None,
 ) -> Tensor:
-    """Convert a NumPy ndarray to an infinicore Tensor.
+    """
+    Creates a Tensor from a numpy.ndarray.
 
     Args:
         np_array: NumPy ndarray to convert to tensor
-        dtype: Optional infinicore dtype. If None, inferred from numpy array
         device: Optional infinicore device. If None, defaults to CPU device
 
     Returns:
         Tensor: An infinicore tensor created from the numpy array
 
     Raises:
         TypeError: If input data is not a numpy ndarray
-        ValueError: If input array is empty
+        ValueError: If input array is empty or not C-contiguous
 
     Note:
         NumPy arrays can only be created on CPU. For CUDA devices, data is first
         created on CPU, then copied to the target device.
     """
-    # Input validation
+
     if not isinstance(np_array, np.ndarray):
         raise TypeError(
             f"Input data must be a np.ndarray, got {type(np_array).__name__}"
@@ -228,55 +263,29 @@ def from_numpy(
     if np_array.size == 0:
         raise ValueError("Input array cannot be empty")
 
-    # Determine target numpy dtype
-    # If dtype is specified, convert it to numpy dtype first
-    if dtype is not None:
-        np_dtype = infinicore_to_numpy_dtype(dtype)
-        # Create a copy with the target dtype if dtype doesn't match
-        # Use copy=True to ensure we don't modify the original array
-        if np_dtype != np_array.dtype:
-            np_array = np_array.astype(np_dtype, copy=True)
-        # Ensure C-contiguous layout
-        elif not np_array.flags.c_contiguous:
-            np_array = np.ascontiguousarray(np_array)
-    else:
-        # Ensure C-contiguous layout
-        if not np_array.flags.c_contiguous:
-            np_array = np.ascontiguousarray(np_array)
-
-    # Infer infinicore dtype if not provided
-    infini_type = (
-        dtype if dtype is not None else numpy_to_infinicore_dtype(np_array.dtype)
-    )
+    if not np_array.flags.c_contiguous:
+        raise ValueError("Input array must be C-contiguous")
 
-    # Default to CPU device if not provided
-    infini_device = device if device is not None else infinicore.device("cpu", 0)
+    infini_type = numpy_to_infinicore_dtype(np_array.dtype)
     cpu_device = infinicore.device("cpu", 0)
 
-    # Create a temporary tensor on CPU using from_blob to reference numpy array
-    # This allows us to copy data without keeping numpy array reference
+    # Create a infinicore.Tensor on CPU using from_blob to reference numpy array
     data_ptr = np_array.ctypes.data_as(ctypes.c_void_p).value
-    temp_tensor = Tensor(
+    infini_tensor = Tensor(
         _infinicore.from_blob(
             data_ptr,
             list(np_array.shape),
             dtype=infini_type._underlying,
             device=cpu_device._underlying,
-        )
+        ),
+        _numpy_ref=np_array,
     )
 
-    # Always create the result tensor on CPU first, then copy data
-    # This ensures we have a proper copy of the data
-    result = empty(list(np_array.shape), dtype=infini_type, device=cpu_device)
-    result.copy_(temp_tensor)
-
     # If target device is not CPU, move the tensor to the target device
-    # The temporary tensor and numpy array will be garbage collected
-    # since we don't keep references to them
-    if infini_device.type != "cpu":
-        result = result.to(infini_device)
+    if device is not None and device.type != "cpu":
+        infini_tensor = infini_tensor.to(device)
 
-    return result
+    return infini_tensor
 
 
 def from_list(data, *, dtype=None, device=None) -> Tensor:
@@ -329,4 +338,4 @@ def from_list(data, *, dtype=None, device=None) -> Tensor:
 
     # Reuse from_numpy to create the tensor
     # This avoids code duplication and ensures consistent behavior
-    return from_numpy(np_array, dtype=dtype, device=device)
+    return from_numpy(np_array, device=device)

python/infinicore/utils.py

@@ -2,6 +2,12 @@
 import numpy as np
 import torch
 
+try:
+    import torch
+except ImportError:
+    torch = None
+    print("warning: torch not available, some functions may not be available")
+
 import infinicore
 
 
@@ -94,4 +100,6 @@ def infinicore_to_numpy_dtype(infini_dtype):
     elif infini_dtype == infinicore.uint8:
         return np.uint8
     else:
-        raise ValueError(f"Unsupported infinicore dtype: {infini_dtype}")
+        raise ValueError(
+            f"Cannot convert infinicore dtype: {infini_dtype} to numpy dtype"
+        )

test/infinicore/test.py

@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 from infinicore.lib import _infinicore
 
 import infinicore
@@ -265,10 +266,161 @@ def func6_initialize_device_relationship():
         z_infini.debug()
 
 
+def test7_infinicore_tensor_function():
+    """
+    测试 infinicore.tensor 函数，能够传入 list, tuple, NumPy, scalar，得到一个InfiniCore.Tensor的对象
+    """
+    print("\n" + "=" * 60)
+    print("测试 infinicore.tensor 函数")
+    print("=" * 60)
+
+    # 定义测试用例列表
+    case_list = [
+        {
+            "name": "从 list 创建 tensor",
+            "data": [[1.0, 2.0, 3.0, 4.0]],
+            "kwargs": {},
+            "expected_shape": [1, 4],
+            "expected_dtype": None,
+            "expected_device_type": None,
+        },
+        {
+            "name": "从 tuple 创建 tensor",
+            "data": (1.0, 2.0, 3.0, 4.0),
+            "kwargs": {},
+            "expected_shape": [4],
+            "expected_dtype": None,
+            "expected_device_type": None,
+        },
+        {
+            "name": "从 NumPy.ndarray 创建 tensor",
+            "data": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+            "kwargs": {},
+            "expected_shape": [2, 2],
+            "expected_dtype": None,
+            "expected_device_type": None,
+        },
+        {
+            "name": "从 scalar (int) 创建 tensor",
+            "data": 42,
+            "kwargs": {},
+            "expected_shape": [],
+            "expected_dtype": None,
+            "expected_device_type": None,
+        },
+        {
+            "name": "从 scalar (float) 创建 tensor",
+            "data": 3.14,
+            "kwargs": {},
+            "expected_shape": [],
+            "expected_dtype": None,
+            "expected_device_type": None,
+        },
+        {
+            "name": "多维 list",
+            "data": [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
+            "kwargs": {},
+            "expected_shape": [3, 2],
+            "expected_dtype": None,
+            "expected_device_type": None,
+        },
+        {
+            "name": "指定 dtype (float32)",
+            "data": [1, 2, 3],
+            "kwargs": {"dtype": infinicore.float32},
+            "expected_shape": [3],
+            "expected_dtype": infinicore.float32,
+            "expected_device_type": None,
+        },
+        {
+            "name": "指定 dtype (float64)",
+            "data": [1, 2, 3],
+            "kwargs": {"dtype": infinicore.float64},
+            "expected_shape": [3],
+            "expected_dtype": infinicore.float64,
+            "expected_device_type": None,
+        },
+        {
+            "name": "指定 device (cuda)",
+            "data": [1.0, 2.0, 3.0],
+            "kwargs": {"device": infinicore.device("cuda", 0)},
+            "expected_shape": [3],
+            "expected_dtype": None,
+            "expected_device_type": "cuda",
+        },
+    ]
+
+    # 循环测试每个用例
+    for i, case in enumerate(case_list, 1):
+        print(f"\n{i}. 测试{case['name']}:")
+        print("-" * 40)
+
+        # 准备输入数据描述
+        if isinstance(case["data"], np.ndarray):
+            input_desc = f"shape={case['data'].shape}, dtype={case['data'].dtype}"
+        else:
+            input_desc = str(case["data"])
+
+        print(f"  输入: {input_desc}")
+
+        # 创建 tensor
+        tensor = infinicore.tensor(case["data"], **case["kwargs"])
+
+        # 打印输出信息
+        print(
+            f"  输出: shape={tensor.shape}, dtype={tensor.dtype}, device={tensor.device}"
+        )
+
+        # 验证结果
+        if case["expected_shape"] is not None:
+            assert tensor.shape == case["expected_shape"], (
+                f"期望shape {case['expected_shape']}, 实际 {tensor.shape}"
+            )
+
+        if case["expected_dtype"] is not None:
+            assert tensor.dtype == case["expected_dtype"], (
+                f"期望dtype {case['expected_dtype']}, 实际 {tensor.dtype}"
+            )
+
+        if case["expected_device_type"] is not None:
+            assert tensor.device.type == case["expected_device_type"], (
+                f"期望device类型 {case['expected_device_type']}, 实际 {tensor.device.type}"
+            )
+
+        print(f"  ✓ {case['name']} 测试通过")
+
+    # 特殊测试：数据正确性验证（与 NumPy/Torch 对比）
+    print("\n10. 测试数据正确性验证（与 NumPy 对比）:")
+    print("-" * 40)
+    test_data = [[1.5, 2.5], [3.5, 4.5]]
+    infini_tensor = infinicore.tensor(test_data, dtype=infinicore.float32)
+
+    # 转换为 torch tensor 进行验证
+    torch_ref = torch.tensor(test_data, dtype=torch.float32)
+    torch_result = torch.zeros(infini_tensor.shape, dtype=torch.float32)
+    infini_blob = infinicore.from_blob(
+        torch_result.data_ptr(),
+        infini_tensor.shape,
+        dtype=infinicore.float32,
+        device=infinicore.device("cpu", 0),
+    )
+    infini_blob.copy_(infini_tensor)
+
+    max_error = torch.abs(torch_ref - torch_result).max().item()
+    print(f"  最大误差: {max_error}")
+    assert max_error < 1e-6, f"数据不匹配，最大误差: {max_error}"
+    print("  ✓ 数据正确性验证通过")
+
+    print("\n" + "=" * 60)
+    print("所有 infinicore.tensor 测试通过！")
+    print("=" * 60 + "\n")
+
+
 if __name__ == "__main__":
     test()
     test2()
     test3()
     test4_to()
     test5_bf16()
     func6_initialize_device_relationship()
+    test7_infinicore_tensor_function()