Improve support for DALI enum types

dali/kernels/common/cast_gpu.cu

@@ -19,6 +19,7 @@
 #include "dali/kernels/common/utils.h"
 #include "dali/kernels/kernel.h"
 #include "dali/kernels/dynamic_scratchpad.h"
+#include "dali/pipeline/data/types.h"
 
 namespace dali {
 namespace kernels {
@@ -102,32 +103,38 @@ void CastGPU<Out, In>::Run(KernelContext &ctx,
 
 #define INSTANTIATE_IMPL(Out, In) template struct DLL_PUBLIC CastGPU<Out, In>;
 
-#define INSTANTIATE_FOREACH_INTYPE(Out)  \
-  INSTANTIATE_IMPL(Out, bool);  \
-  INSTANTIATE_IMPL(Out, uint8_t);  \
-  INSTANTIATE_IMPL(Out, uint16_t); \
-  INSTANTIATE_IMPL(Out, uint32_t);  \
-  INSTANTIATE_IMPL(Out, uint64_t);  \
-  INSTANTIATE_IMPL(Out, int8_t);  \
-  INSTANTIATE_IMPL(Out, int16_t);  \
-  INSTANTIATE_IMPL(Out, int32_t);  \
-  INSTANTIATE_IMPL(Out, int64_t);  \
-  INSTANTIATE_IMPL(Out, float);    \
-  INSTANTIATE_IMPL(Out, double);    \
-  INSTANTIATE_IMPL(Out, dali::float16);
-
-INSTANTIATE_FOREACH_INTYPE(bool);  \
-INSTANTIATE_FOREACH_INTYPE(uint8_t);  \
-INSTANTIATE_FOREACH_INTYPE(uint16_t); \
-INSTANTIATE_FOREACH_INTYPE(uint32_t);  \
-INSTANTIATE_FOREACH_INTYPE(uint64_t);  \
-INSTANTIATE_FOREACH_INTYPE(int8_t);  \
-INSTANTIATE_FOREACH_INTYPE(int16_t);  \
-INSTANTIATE_FOREACH_INTYPE(int32_t);  \
-INSTANTIATE_FOREACH_INTYPE(int64_t);  \
-INSTANTIATE_FOREACH_INTYPE(float);    \
-INSTANTIATE_FOREACH_INTYPE(double);    \
+#define INSTANTIATE_FOREACH_INTYPE(Out) \
+  INSTANTIATE_IMPL(Out, bool);          \
+  INSTANTIATE_IMPL(Out, uint8_t);       \
+  INSTANTIATE_IMPL(Out, uint16_t);      \
+  INSTANTIATE_IMPL(Out, uint32_t);      \
+  INSTANTIATE_IMPL(Out, uint64_t);      \
+  INSTANTIATE_IMPL(Out, int8_t);        \
+  INSTANTIATE_IMPL(Out, int16_t);       \
+  INSTANTIATE_IMPL(Out, int32_t);       \
+  INSTANTIATE_IMPL(Out, int64_t);       \
+  INSTANTIATE_IMPL(Out, float);         \
+  INSTANTIATE_IMPL(Out, double);        \
+  INSTANTIATE_IMPL(Out, dali::float16); \
+  INSTANTIATE_IMPL(Out, DALIDataType);  \
+  INSTANTIATE_IMPL(Out, DALIImageType); \
+  INSTANTIATE_IMPL(Out, DALIInterpType);
+
+INSTANTIATE_FOREACH_INTYPE(bool);
+INSTANTIATE_FOREACH_INTYPE(uint8_t);
+INSTANTIATE_FOREACH_INTYPE(uint16_t);
+INSTANTIATE_FOREACH_INTYPE(uint32_t);
+INSTANTIATE_FOREACH_INTYPE(uint64_t);
+INSTANTIATE_FOREACH_INTYPE(int8_t);
+INSTANTIATE_FOREACH_INTYPE(int16_t);
+INSTANTIATE_FOREACH_INTYPE(int32_t);
+INSTANTIATE_FOREACH_INTYPE(int64_t);
+INSTANTIATE_FOREACH_INTYPE(float);
+INSTANTIATE_FOREACH_INTYPE(double);
 INSTANTIATE_FOREACH_INTYPE(dali::float16);
+INSTANTIATE_FOREACH_INTYPE(DALIDataType);
+INSTANTIATE_FOREACH_INTYPE(DALIImageType);
+INSTANTIATE_FOREACH_INTYPE(DALIInterpType);
 
 }  // namespace cast
 }  // namespace kernels

dali/operators/generic/cast.h

@@ -19,13 +19,14 @@
 
 #include "dali/core/convert.h"
 #include "dali/core/tensor_shape.h"
+#include "dali/pipeline/data/types.h"
 #include "dali/pipeline/operator/checkpointing/stateless_operator.h"
 
 namespace dali {
 
 #define CAST_ALLOWED_TYPES                                                                         \
   (bool, uint8_t, uint16_t, uint32_t, uint64_t, int8_t, int16_t, int32_t, int64_t, float16, float, \
-  double)
+  double, DALIDataType, DALIImageType, DALIInterpType)
 
 template <typename Backend>
 class Cast : public StatelessOperator<Backend> {
@@ -54,6 +55,11 @@ class Cast : public StatelessOperator<Backend> {
   bool SetupImpl(std::vector<OutputDesc> &output_desc, const Workspace &ws) override {
     const auto &input = ws.Input<Backend>(0);
     DALIDataType out_type = is_cast_like_ ?  ws.GetInputDataType(1) : dtype_arg_;
+    DALI_ENFORCE(!(IsEnum(input.type()) && IsFloatingPoint(out_type) ||
+                   IsEnum(out_type) && IsFloatingPoint(input.type())),
+                 make_string("Cannot cast from ", input.type(), " to ", out_type,
+                             ". Enums can only participate in casts with integral types, "
+                             "but not floating point types."));
     output_desc.resize(1);
     output_desc[0].shape = input.shape();
     output_desc[0].type = out_type;

dali/operators/generic/constant.h

@@ -22,8 +22,9 @@
 #include "dali/core/tensor_view.h"
 #include "dali/core/static_switch.h"
 
-#define CONSTANT_OP_SUPPORTED_TYPES \
-  (bool, int8_t, uint8_t, int16_t, uint16_t, int32_t, uint32_t, int64_t, uint64_t, float, float16)
+#define CONSTANT_OP_SUPPORTED_TYPES                                                                \
+  (bool, int8_t, uint8_t, int16_t, uint16_t, int32_t, uint32_t, int64_t, uint64_t, float, float16, \
+  DALIDataType, DALIImageType, DALIInterpType)
 
 namespace dali {
 

dali/operators/random/choice.h

@@ -29,7 +29,7 @@
   uint8_t, uint16_t, uint32_t, uint64_t, int8_t, int16_t, int32_t, int64_t
 #define DALI_CHOICE_1D_TYPES                                                                      \
   bool, uint8_t, uint16_t, uint32_t, uint64_t, int8_t, int16_t, int32_t, int64_t, float16, float, \
-      double
+      double, DALIDataType, DALIImageType, DALIInterpType
 
 namespace dali {
 

dali/operators/random/choice_cpu.cc

@@ -31,6 +31,9 @@ a single value per sample is generated.
 
 The type of the output matches the type of the input.
 For scalar inputs, only integral types are supported, otherwise any type can be used.
+The operator supports selection from an input containing elements of one of DALI enum types,
+that is: :meth:`nvidia.dali.types.DALIDataType`, :meth:`nvidia.dali.types.DALIImageType`, or
+:meth:`nvidia.dali.types.DALIInterpType`.
 )code")
     .NumInput(1, 2)
     .InputDox(0, "a", "scalar or TensorList",

dali/pipeline/data/types.h

@@ -277,6 +277,18 @@ constexpr bool IsUnsigned(DALIDataType type) {
   }
 }
 
+
+constexpr bool IsEnum(DALIDataType type) {
+  switch (type) {
+    case DALI_DATA_TYPE:
+    case DALI_IMAGE_TYPE:
+    case DALI_INTERP_TYPE:
+      return true;
+    default:
+      return false;
+  }
+}
+
 template <DALIDataType id>
 struct id2type_helper;
 

dali/python/backend_impl.cc

@@ -627,7 +627,7 @@ void ExposeTensor(py::module &m) {
       return FromPythonTrampoline("nvidia.dali.tensors", "_tensor_to_string")(t);
     })
     .def("__repr__", [](Tensor<CPUBackend> &t) {
-      return FromPythonTrampoline("nvidia.dali.tensors", "_tensor_to_string")(t);
+      return FromPythonTrampoline("nvidia.dali.tensors", "_tensor_to_string")(t, false);
     })
     .def_property("__array_interface__", &ArrayInterfaceRepr<CPUBackend>, nullptr,
       R"code(
@@ -777,7 +777,7 @@ void ExposeTensor(py::module &m) {
       return FromPythonTrampoline("nvidia.dali.tensors", "_tensor_to_string")(t);
     })
     .def("__repr__", [](Tensor<GPUBackend> &t) {
-      return FromPythonTrampoline("nvidia.dali.tensors", "_tensor_to_string")(t);
+      return FromPythonTrampoline("nvidia.dali.tensors", "_tensor_to_string")(t, false);
     })
     .def_property("__cuda_array_interface__",  &ArrayInterfaceRepr<GPUBackend>, nullptr,
       R"code(
@@ -1193,7 +1193,11 @@ void ExposeTensorList(py::module &m) {
       return FromPythonTrampoline("nvidia.dali.tensors", "_tensorlist_to_string")(t);
     })
     .def("__repr__", [](TensorList<CPUBackend> &t) {
-      return FromPythonTrampoline("nvidia.dali.tensors", "_tensorlist_to_string")(t);
+      // Repr might be used in exceptions and the data might not be possible to be represented
+      // (DALI enums do not support buffer protocol due to difference between C++ numeric
+      // representation and Python "O" - object/pointer-based representation).
+      // That why we skip the data part.
+      return FromPythonTrampoline("nvidia.dali.tensors", "_tensorlist_to_string")(t, false);
     })
     .def_property_readonly("dtype", [](TensorList<CPUBackend> &tl) {
           return tl.type();
@@ -1393,7 +1397,7 @@ void ExposeTensorList(py::module &m) {
       return FromPythonTrampoline("nvidia.dali.tensors", "_tensorlist_to_string")(t);
     })
     .def("__repr__", [](TensorList<GPUBackend> &t) {
-      return FromPythonTrampoline("nvidia.dali.tensors", "_tensorlist_to_string")(t);
+      return FromPythonTrampoline("nvidia.dali.tensors", "_tensorlist_to_string")(t, false);
     })
     .def_property_readonly("dtype", [](TensorList<GPUBackend> &tl) {
           return tl.type();

dali/python/nvidia/dali/_backend_enums.pyi

@@ -29,6 +29,9 @@ class DALIDataType(Enum):
     FLOAT64 = ...
     BOOL = ...
     STRING = ...
+    DATA_TYPE = ...
+    IMAGE_TYPE = ...
+    INTERP_TYPE = ...
 
 class DALIImageType(Enum):
     RGB = ...

dali/python/nvidia/dali/_debug_mode.py

@@ -45,7 +45,7 @@ def __str__(self):
         indent = " " * 4
         return (
             f'DataNodeDebug(\n{indent}name="{self.name}",\n{indent}data='
-            + f'{_tensors._tensorlist_to_string(self._data, indent + " " * 5)})'
+            + f'{_tensors._tensorlist_to_string(self._data, indent=indent + " " * 5)})'
         )
 
     __repr__ = __str__

dali/python/nvidia/dali/tensors.py

@@ -50,35 +50,56 @@ def import_numpy():
             )
 
 
-def _tensor_to_string(self):
-    """Returns string representation of Tensor."""
+def _tensor_to_string(self, show_data=True):
+    """Returns string representation of Tensor.
+
+    Parameters
+    ----------
+    show_data : bool, optional
+        Access and format the underlying data, by default True
+    """
     import_numpy()
 
     type_name = type(self).__name__
     indent = " " * 4
     layout = self.layout()
-    data = np.array(_transfer_to_cpu(self, type_name[-3:]))
-    data_str = np.array2string(data, prefix=indent, edgeitems=2)
+    if show_data:
+        data = np.array(_transfer_to_cpu(self, type_name[-3:]))
+        data_str = np.array2string(data, prefix=indent, edgeitems=2)
 
     params = (
-        [f"{type_name}(\n{indent}{data_str}", f"dtype={self.dtype}"]
+        ([f"{data_str}"] if show_data else [])
+        + [f"dtype={self.dtype}"]
         + ([f"layout={layout}"] if layout else [])
         + [f"shape={self.shape()})"]
     )
 
-    return _join_string(params, False, 0, ",\n" + indent)
+    return f"{type_name}(\n{indent}" + _join_string(params, False, 0, ",\n" + indent)
 
 
-def _tensorlist_to_string(self, indent=""):
-    """Returns string representation of TensorList."""
+def _tensorlist_to_string(self, show_data=True, indent=""):
+    """Returns string representation of TensorList.
+
+    Parameters
+    ----------
+    show_data : bool, optional
+        Access and format the underlying data, by default True
+    indent : str, optional
+        optional indentation used in formatting, by default ""
+    """
     import_numpy()
 
     edgeitems = 2
     spaces_indent = indent + " " * 4
     type_name = type(self).__name__
     layout = self.layout()
-    data = _transfer_to_cpu(self, type_name[-3:])
-    data_str = "[]"
+    if show_data:
+        data = _transfer_to_cpu(self, type_name[-3:])
+        data_str = "[]"
+    else:
+        data = None
+        data_str = ""
+
     crop = False
 
     if data:
@@ -120,9 +141,10 @@ def _tensorlist_to_string(self, indent=""):
         )
 
     params = (
-        [f"{type_name}(\n{spaces_indent}{data_str}", f"dtype={self.dtype}"]
+        ([f"{data_str}"] if show_data else [])
+        + [f"dtype={self.dtype}"]
         + ([f'layout="{layout}"'] if layout else [])
         + [f"num_samples={len(self)}", f"{shape_prefix}{shape_str}])"]
     )
 
-    return _join_string(params, False, 0, ",\n" + spaces_indent)
+    return f"{type_name}(\n{spaces_indent}" + _join_string(params, False, 0, ",\n" + spaces_indent)

dali/python/nvidia/dali/types.py

@@ -519,16 +519,45 @@ def _type_from_value_or_list(v):
             has_floats = False
             has_ints = False
             has_bools = False
+            has_enums = False
+            enum_type = None
             for x in v:
                 if isinstance(x, float):
                     has_floats = True
                 elif isinstance(x, bool):
                     has_bools = True
                 elif isinstance(x, int):
                     has_ints = True
+                elif isinstance(x, (DALIDataType, DALIImageType, DALIInterpType)):
+                    has_enums = True
+                    enum_type = type(x)
+                    break
                 else:
                     raise TypeError("Unexpected type: " + str(type(x)))
 
+            if has_enums:
+                for x in v:
+                    if not isinstance(x, enum_type):
+                        raise TypeError(
+                            f"Expected all elements of the input to be the "
+                            f"same enum type: `{enum_type.__name__}` but got `{type(x).__name__}` "
+                            f"for one of the elements."
+                        )
+
+            if has_enums:
+                if issubclass(enum_type, DALIDataType):
+                    return DALIDataType.DATA_TYPE
+                elif issubclass(enum_type, DALIImageType):
+                    return DALIDataType.IMAGE_TYPE
+                elif issubclass(enum_type, DALIInterpType):
+                    return DALIDataType.INTERP_TYPE
+                else:
+                    raise TypeError(
+                        f"Unexpected enum type: `{enum_type.__name__}`, expected one of: "
+                        "`nvidia.dali.types.DALIDataType`, `nvidia.dali.types.DALIImageType`, "
+                        "or `nvidia.dali.types.DALIInterpType`."
+                    )
+
             if has_floats:
                 return DALIDataType.FLOAT
             if has_ints:
@@ -582,7 +611,8 @@ def Constant(value, dtype=None, shape=None, layout=None, device=None, **kwargs):
 
     Args
     ----
-    value: `bool`, `int`, `float`, a `list` or `tuple` thereof or a `numpy.ndarray`
+    value: `bool`, `int`, `float`, `DALIDataType` `DALIImageType`, `DALIInterpType`,
+           a `list` or `tuple` thereof or a `numpy.ndarray`
         The constant value to wrap. If it is a scalar, it can be used as scalar
         value in mathematical expressions. Otherwise, it will produce a constant
         tensor node (optionally reshaped according to `shape` argument).
@@ -606,10 +636,24 @@ def Constant(value, dtype=None, shape=None, layout=None, device=None, **kwargs):
         and the arguments are passed to the `dali.ops.Constant` operator
     """
 
+    def is_enum(value, dtype):
+        # we force true scalar enums through a Constant node rather than using ScalarConstant
+        # as they do not support any arithmetic operations
+        if isinstance(value, (DALIDataType, DALIImageType, DALIInterpType)):
+            return True
+        elif dtype is not None and dtype in {
+            DALIDataType.DATA_TYPE,
+            DALIDataType.IMAGE_TYPE,
+            DALIDataType.INTERP_TYPE,
+        }:
+            return True
+        return False
+
     if (
         device is not None
         or (_is_compatible_array_type(value) and not _is_true_scalar(value))
         or isinstance(value, (list, tuple))
+        or is_enum(value, dtype)
         or not _is_scalar_shape(shape)
         or kwargs
         or layout is not None