From ef603eb546543243fc0577a1e6d823427e8aafe6 Mon Sep 17 00:00:00 2001 From: pengcheng888 Date: Tue, 13 Jan 2026 14:31:56 +0800 Subject: [PATCH] =?UTF-8?q?issue/892-=E6=B7=BB=E5=8A=A0infini.tensor?= =?UTF-8?q?=E5=87=BD=E6=95=B0;=20=E7=A7=BB=E9=99=A4from=5Fnumpy=E5=87=BD?= =?UTF-8?q?=E6=95=B0=E7=9A=84dtype=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/infinicore/__init__.py | 2 + python/infinicore/nn/modules/rope.py | 10 +- python/infinicore/tensor.py | 95 +++++++++-------- python/infinicore/utils.py | 10 +- test/infinicore/test.py | 152 +++++++++++++++++++++++++++ 5 files changed, 223 insertions(+), 46 deletions(-) diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index b7288f3ac..2fafc3a85 100644 --- a/python/infinicore/__init__.py +++ b/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", ] diff --git a/python/infinicore/nn/modules/rope.py b/python/infinicore/nn/modules/rope.py index 5d579e2d3..1a53099f6 100644 --- a/python/infinicore/nn/modules/rope.py +++ b/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 diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py index 8e2c9b2d6..c75041b7f 100644 --- a/python/infinicore/tensor.py +++ b/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,14 +234,13 @@ 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: @@ -213,13 +248,13 @@ def from_numpy( 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) diff --git a/python/infinicore/utils.py b/python/infinicore/utils.py index 094b2230e..20b515781 100644 --- a/python/infinicore/utils.py +++ b/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" + ) diff --git a/test/infinicore/test.py b/test/infinicore/test.py index 36aeffe4e..d8946cbde 100644 --- a/test/infinicore/test.py +++ b/test/infinicore/test.py @@ -1,4 +1,5 @@ import torch +import numpy as np from infinicore.lib import _infinicore import infinicore @@ -265,6 +266,156 @@ 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() @@ -272,3 +423,4 @@ def func6_initialize_device_relationship(): test4_to() test5_bf16() func6_initialize_device_relationship() + test7_infinicore_tensor_function()