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..d5917ac0e 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) + 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,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..9ba0e0874 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,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()