diff --git a/dali/python/nvidia/dali/experimental/dynamic/_batch.py b/dali/python/nvidia/dali/experimental/dynamic/_batch.py index 693838390f6..d9a910418a9 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_batch.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_batch.py @@ -17,7 +17,6 @@ from collections.abc import Iterator, Sequence import nvidia.dali.backend as _backend -import nvidia.dali.types as _dali_types import nvidia.dali._tensor_formatting as _tensor_formatting from ._nvtx import NVTXRange from nvidia.dali._typing import BatchLike, TensorLike @@ -29,12 +28,12 @@ from ._arithmetic import _arithm_op from ._device import Device, DeviceLike from ._device import device as _device -from ._tensor import Tensor, _is_full_slice, _try_convert_enums +from ._tensor import Tensor, _array_from_python, _is_full_slice from ._tensor import as_tensor as _as_tensor from ._tensor import tensor as _tensor from ._type import DType, DTypeLike from ._type import dtype as _dtype -from .capture._invariant import unwrap_invariant_args, unwrap_invariants +from .capture._invariant import unwrap_invariant_args def _backend_device(backend: _backend.TensorListCPU | _backend.TensorListGPU) -> Device: @@ -435,10 +434,6 @@ def assign_fast_path_storage(storage, dev, wraps_external_data): def _is_external(self) -> bool: return self._wraps_external_data - _nvtx_to_numpy_and_stack = NVTXRange("broadcast: to numpy and stack", category="batch") - _nvtx_to_backend = NVTXRange("broadcast: to backend", category="batch") - _nvtx_create_batch = NVTXRange("broadcast: create batch", category="batch") - @staticmethod @NVTXRange("broadcast", category="batch") def broadcast( @@ -461,36 +456,21 @@ def broadcast( sample, batch_size, device, dtype = unwrap_invariant_args(sample, batch_size, device, dtype) if isinstance(sample, Batch): raise ValueError("Cannot broadcast a Batch") - if _is_tensor_type(sample): - t = _as_tensor(sample, device=device, dtype=dtype).evaluate() - if t.device.device_type == "gpu": - tl_type = _backend.TensorListGPU - else: - tl_type = _backend.TensorListCPU - return Batch(tl_type.broadcast(t._storage, batch_size)) - import numpy as np - - with Batch._nvtx_to_numpy_and_stack: - arr = np.array(unwrap_invariants(sample)) - converted_dtype_id = None - if arr.dtype == np.float64: - arr = arr.astype(np.float32) - elif arr.dtype == np.int64: - arr = arr.astype(np.int32) - elif arr.dtype == np.uint64: - arr = arr.astype(np.uint32) - elif arr.dtype == object: - arr, converted_dtype_id = _try_convert_enums(arr) - if dtype is not None and dtype.kind != DType.Kind.enum: - arr = arr.astype(_dali_types.to_numpy_type(dtype.type_id)) - arr = np.repeat(arr[np.newaxis], batch_size, axis=0) + if not _is_tensor_type(sample): + import numpy as np - with Batch._nvtx_to_backend: + arr, converted_dtype_id = _array_from_python(sample, dtype) + # Materialize Python constants contiguously for bulk CPU/GPU transfers. + arr = np.repeat(arr[np.newaxis], batch_size, axis=0) tl = _backend.TensorListCPU(arr) if converted_dtype_id is not None: tl.reinterpret(converted_dtype_id) - with Batch._nvtx_create_batch: - return Batch(tl, device=device, dtype=dtype) + return Batch(tl, device=device) + t = _as_tensor(sample, device=device, dtype=dtype).evaluate() + tl_type = ( + _backend.TensorListGPU if t.device.device_type == "gpu" else _backend.TensorListCPU + ) + return Batch(tl_type.broadcast(t._storage, batch_size)) @property def dtype(self) -> DType: diff --git a/dali/python/nvidia/dali/experimental/dynamic/_tensor.py b/dali/python/nvidia/dali/experimental/dynamic/_tensor.py index 53dc1a46444..e7ca12049f6 100644 --- a/dali/python/nvidia/dali/experimental/dynamic/_tensor.py +++ b/dali/python/nvidia/dali/experimental/dynamic/_tensor.py @@ -78,6 +78,52 @@ def _try_convert_enums(arr): raise TypeError(f"Unexpected element type {type(item)}") +def _is_wide_int_array(arr: np.ndarray): + if arr.dtype in (np.int64, np.uint64): + return True + return arr.dtype == object and arr.size > 0 and all(type(value) is int for value in arr.flat) + + +def _convert_integer_array(arr: np.ndarray): + """Infer int32 or uint32 for integer data and prevent overflow.""" + dtype = np.int32 + if arr.size: + min_value, max_value = int(arr.min()), int(arr.max()) + if min_value >= 0 and max_value >> 31: + dtype = np.uint32 + limits = np.iinfo(dtype) + if min_value < limits.min or max_value > limits.max: + value = min_value if min_value < limits.min else max_value + raise OverflowError(f"Python integer {value} out of bounds for {limits.dtype}.") + return arr.astype(dtype, copy=False) + + +def _array_from_python(data, dtype=None): + """Convert Python data to NumPy, returning any DALI enum type to reinterpret.""" + data = unwrap_invariants(data) + converted_dtype_id = None + if dtype is not None: + if not isinstance(dtype, DType): + dtype = _dtype(dtype) + if dtype.kind == DType.Kind.enum: + numpy_type = np.int32 + converted_dtype_id = dtype.type_id + else: + numpy_type = nvidia.dali.types.to_numpy_type(dtype.type_id) + + arr = np.array(data, dtype=numpy_type) + else: + arr = np.array(data) + # Infer 32-bit types for Python numbers, preserving integer values. + if _is_wide_int_array(arr): + arr = _convert_integer_array(arr) + elif arr.dtype == np.float64: + arr = arr.astype(np.float32) + elif arr.dtype == object: + arr, converted_dtype_id = _try_convert_enums(arr) + return arr, converted_dtype_id + + class Tensor: """A Tensor object. @@ -260,45 +306,16 @@ def __init__( else: raise ValueError(f"Unsupported device type: {dl_device_type}") self._wraps_external_data = True - elif a := _get_array_interface(data): + elif (a := _get_array_interface(data)) is not None: self._storage = _backend.TensorCPU(a, layout) self._wraps_external_data = True else: - if dtype is not None: - if dtype.kind == DType.Kind.enum: - numpy_type = np.int32 - else: - numpy_type = nvidia.dali.types.to_numpy_type(dtype.type_id) - - self._storage = _backend.TensorCPU( - np.array(unwrap_invariants(data), dtype=numpy_type), - layout, - False, - ) - if dtype.kind == DType.Kind.enum: - self._storage.reinterpret(dtype.type_id) - - copied = True - self._wraps_external_data = False - self._dtype = dtype - else: - arr = np.array(unwrap_invariants(data)) - # DALI doesn't support int64 and float64, so we need to convert them to int32 - # and float32, respectively. - converted_dtype_id = None - if arr.dtype == np.int64: - arr = arr.astype(np.int32) - elif arr.dtype == np.uint64: - arr = arr.astype(np.uint32) - elif arr.dtype == np.float64: - arr = arr.astype(np.float32) - elif arr.dtype == object: - arr, converted_dtype_id = _try_convert_enums(arr) - self._storage = _backend.TensorCPU(arr, layout, False) - if converted_dtype_id is not None: - self._storage.reinterpret(converted_dtype_id) - copied = True - self._wraps_external_data = False + arr, converted_dtype_id = _array_from_python(data, dtype) + self._storage = _backend.TensorCPU(arr, layout, False) + if converted_dtype_id is not None: + self._storage.reinterpret(converted_dtype_id) + copied = True + self._wraps_external_data = False if self._storage is not None: self._device = _backend_device(self._storage) diff --git a/dali/test/python/experimental_mode/test_batch.py b/dali/test/python/experimental_mode/test_batch.py index e33d987d45b..0d117e996d3 100644 --- a/dali/test/python/experimental_mode/test_batch.py +++ b/dali/test/python/experimental_mode/test_batch.py @@ -128,6 +128,21 @@ def test_broadcast(): assert np.array_equal(np.array(t), a) +@eval_modes() +@params(((1 << 31) - 1, ndd.int32), (1 << 31, ndd.uint32)) +def test_broadcast_int_dtype(value, dtype): + batch = ndd.Batch.broadcast(value, 5) + assert batch.dtype == dtype + assert all(sample.item() == value for sample in batch.tensors) + + +@eval_modes() +@params((1 << 32,), ([0, 5_000_000_000])) +def test_broadcast_int_overflow(data): + with assert_raises(OverflowError, glob=f"*{data}*out of range for uint32*"): + ndd.Batch.broadcast(data, 5) + + def batch_equal(a, b): if len(a) != len(b): return False diff --git a/dali/test/python/experimental_mode/test_tensor.py b/dali/test/python/experimental_mode/test_tensor.py index 0b96dd597f3..305bd5d33bc 100644 --- a/dali/test/python/experimental_mode/test_tensor.py +++ b/dali/test/python/experimental_mode/test_tensor.py @@ -201,6 +201,55 @@ def test_scalar(): assert scalar._storage.dtype == ndd.int32.type_id +@eval_modes() +@params( + (-(1 << 31), ndd.int32), + ((1 << 31) - 1, ndd.int32), + (1 << 31, ndd.uint32), + ((1 << 32) - 1, ndd.uint32), + ([-(1 << 31), 0, (1 << 31) - 1], ndd.int32), + ([0, 1 << 31, (1 << 32) - 1], ndd.uint32), +) +def test_int_dtype(data, dtype): + tensor = ndd.tensor(data) + assert tensor.dtype == dtype + assert np.array_equal(tensor, data) + + +@eval_modes() +@params( + (-(1 << 31) - 1, -(1 << 31) - 1, "int32"), + (1 << 32, 1 << 32, "uint32"), + ([0, 1 << 80], 1 << 80, "uint32"), + ([-1, 1 << 31], 1 << 31, "int32"), +) +def test_int_overflow(data, faulty_value, dtype): + with assert_raises(OverflowError, glob=f"{faulty_value}*out of range for {dtype}"): + ndd.tensor(data) + + +@eval_modes() +@params( + (1 << 31, np.int32), + ([np.int64(1 << 31)], np.int32), + ([np.uint64(1 << 32)], np.uint32), + ([np.int32(-1)], np.uint32), + (5_000_000_000, np.int64), +) +def test_explicit_int_cast_matches_numpy(data, numpy_type): + dtype = ndd.dtype(numpy_type) + # NumPy may reject Python integers while wrapping NumPy integer scalars. + try: + expected = np.array(data, dtype=numpy_type) + except OverflowError: + with assert_raises(OverflowError): + ndd.tensor(data, dtype=dtype) + else: + tensor = ndd.tensor(data, dtype=dtype) + assert tensor.dtype == dtype + assert np.array_equal(tensor, expected) + + @eval_modes() def test_shapes(): shapes = [