Fix int downcasting in dynamic mode

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:

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)

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

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 = [