Rebase on reinterpret.

Signed-off-by: Michal Zientkiewicz <[email protected]>

Author: mzient

dali/python/nvidia/dali/experimental/dali2/_batch.py

@@ -14,7 +14,7 @@
 
 from typing import Any, Optional, Tuple, List, Union
 from ._type import DType, dtype as _dtype, type_id as _type_id
-from ._tensor import Tensor, _is_full_slice, _is_tensor_type
+from ._tensor import Tensor, _is_full_slice, _is_tensor_type, _try_convert_enums
 import nvidia.dali.backend as _backend
 from ._eval_context import EvalContext as _EvalContext
 from ._device import Device
@@ -164,16 +164,21 @@ def broadcast(sample, batch_size: int, device: Optional[Device] = None) -> "Batc
         import numpy as np
         with nvtx.annotate("to numpy and stack", domain="batch"):
             arr = np.array(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)
             arr = np.repeat(arr[np.newaxis], batch_size, axis=0)
 
         with nvtx.annotate("to backend", domain="batch"):
             tl = _backend.TensorListCPU(arr)
+            if converted_dtype_id is not None:
+                tl.reinterpret(converted_dtype_id)
         with nvtx.annotate("create batch", domain="batch"):
             return Batch(tl, device=device)
 

dali/python/nvidia/dali/experimental/dali2/_op_builder.py

@@ -95,7 +95,7 @@ def _get_batch_size(x):
             return x.batch_size
         if isinstance(x, (_b.TensorListCPU, _b.TensorListGPU)):
             return len(x)
-        if isinstance(x, list) and any(_is_tensor_type(t, True) for t in x):
+        if isinstance(x, list) and any(isinstance(t, Tensor) for t in x):
             return len(x)
         return None
 

dali/python/nvidia/dali/experimental/dali2/_tensor.py

@@ -35,7 +35,7 @@ def _is_tensor_type(x, nested_list_warning=False):
     from . import _batch
 
     if isinstance(x, _batch.Batch):
-        raise ValueError("A list of Batchs is not a valid argument type")
+        raise ValueError("A list of Batch objects is not a valid argument type")
     if isinstance(x, Tensor):
         return True
     if hasattr(x, "__array__"):
@@ -177,7 +177,7 @@ def __init__(
                         (arr, converted_dtype_id) = _try_convert_enums(arr)
                     self._backend = _backend.TensorCPU(arr, layout, False)
                     if converted_dtype_id is not None:
-                        self._backend.reinterpret_as(converted_dtype_id)
+                        self._backend.reinterpret(converted_dtype_id)
                     copied = True
                     self._wraps_external_data = False