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

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)
+    else:
+        raise ValueError(f"Unsupported data type: {type(data)}")
+
+
 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,130 @@ 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)
+
+    # 1. 测试从 list 创建 tensor
+    print("\n1. 测试从 list 创建 tensor:")
+    print("-" * 40)
+    list_data = [[1.0, 2.0, 3.0, 4.0]]
+    tensor_from_list = infinicore.tensor(list_data)
+    print(f"  输入: {list_data}")
+    print(f"  输出: shape={tensor_from_list.shape}, dtype={tensor_from_list.dtype}, device={tensor_from_list.device}")
+    assert tensor_from_list.shape == [1,4], f"期望shape [1,4], 实际 {tensor_from_list.shape}"
+    print("  ✓ list 测试通过")
+
+    # 2. 测试从 tuple 创建 tensor
+    print("\n2. 测试从 tuple 创建 tensor:")
+    print("-" * 40)
+    tuple_data = (1.0, 2.0, 3.0, 4.0)
+    tensor_from_tuple = infinicore.tensor(tuple_data)
+    print(f"  输入: {tuple_data}")
+    print(f"  输出: shape={tensor_from_tuple.shape}, dtype={tensor_from_tuple.dtype}, device={tensor_from_tuple.device}")
+    assert tensor_from_tuple.shape == [4], f"期望shape [4], 实际 {tensor_from_tuple.shape}"
+    print("  ✓ tuple 测试通过")
+
+    # 3. 测试从 NumPy array 创建 tensor
+    print("\n3. 测试从 NumPy array 创建 tensor:")
+    print("-" * 40)
+    np_data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+    tensor_from_numpy = infinicore.tensor(np_data)
+    print(f"  输入: shape={np_data.shape}, dtype={np_data.dtype}")
+    print(f"  输出: shape={tensor_from_numpy.shape}, dtype={tensor_from_numpy.dtype}, device={tensor_from_numpy.device}")
+    assert tensor_from_numpy.shape == [2, 2], f"期望shape [2, 2], 实际 {tensor_from_numpy.shape}"
+    print("  ✓ NumPy array 测试通过")
+
+    # 4. 测试从 scalar (int) 创建 tensor
+    print("\n4. 测试从 scalar (int) 创建 tensor:")
+    print("-" * 40)
+    int_scalar = 42
+    tensor_from_int = infinicore.tensor(int_scalar)
+    print(f"  输入: {int_scalar} (int)")
+    print(f"  输出: shape={tensor_from_int.shape}, dtype={tensor_from_int.dtype}, device={tensor_from_int.device}")
+    assert tensor_from_int.shape == [1], f"期望shape [1], 实际 {tensor_from_int.shape}"
+    print("  ✓ scalar (int) 测试通过")
+
+    # 5. 测试从 scalar (float) 创建 tensor
+    print("\n5. 测试从 scalar (float) 创建 tensor:")
+    print("-" * 40)
+    float_scalar = 3.14
+    tensor_from_float = infinicore.tensor(float_scalar)
+    print(f"  输入: {float_scalar} (float)")
+    print(f"  输出: shape={tensor_from_float.shape}, dtype={tensor_from_float.dtype}, device={tensor_from_float.device}")
+    assert tensor_from_float.shape == [1], f"期望shape [1], 实际 {tensor_from_float.shape}"
+    print("  ✓ scalar (float) 测试通过")
+
+    # 6. 测试指定 dtype
+    print("\n6. 测试指定 dtype:")
+    print("-" * 40)
+    list_data_f32 = [1, 2, 3]
+    tensor_f32 = infinicore.tensor(list_data_f32, dtype=infinicore.float32)
+    tensor_f64 = infinicore.tensor(list_data_f32, dtype=infinicore.float64)
+    print(f"  输入: {list_data_f32}")
+    print(f"  float32: dtype={tensor_f32.dtype}")
+    print(f"  float64: dtype={tensor_f64.dtype}")
+    assert tensor_f32.dtype == infinicore.float32, f"期望 float32, 实际 {tensor_f32.dtype}"
+    assert tensor_f64.dtype == infinicore.float64, f"期望 float64, 实际 {tensor_f64.dtype}"
+    print("  ✓ dtype 指定测试通过")
+
+    # 7. 测试指定 device
+    print("\n7. 测试指定 device:")
+    print("-" * 40)
+    list_data_device = [1.0, 2.0, 3.0]
+    tensor_cpu = infinicore.tensor(list_data_device, device=infinicore.device("cuda", 0))
+    print(f"  输入: {list_data_device}")
+    print(f"  CPU: device={tensor_cpu.device}")
+    assert tensor_cpu.device.type == "cuda", f"期望 CPU, 实际 {tensor_cpu.device.type}"
+    print("  ✓ device 指定测试通过 (cuda)")
+
+    # 8. 测试多维 list
+    print("\n8. 测试多维 list:")
+    print("-" * 40)
+    nested_list = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]
+    tensor_nested = infinicore.tensor(nested_list)
+    print(f"  输入: {nested_list}")
+    print(f"  输出: shape={tensor_nested.shape}, dtype={tensor_nested.dtype}")
+    assert tensor_nested.shape == [3, 2], f"期望shape [3, 2], 实际 {tensor_nested.shape}"
+    print("  ✓ 多维 list 测试通过")
+
+    # 9. 测试数据正确性验证（与 NumPy 对比）
+    print("\n9. 测试数据正确性验证（与 NumPy 对比）:")
+    print("-" * 40)
+    test_data = [[1.5, 2.5], [3.5, 4.5]]
+    np_ref = np.array(test_data, dtype=np.float32)
+    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()
+    # test()
+    # test2()
+    # test3()
+    # test4_to()
+    # test5_bf16()
+    # func6_initialize_device_relationship()
+    test7_infinicore_tensor_function()