[onert] Refactor tensorinfo to be more pythonic and robust

runtime/onert/api/python/include/nnfw_api_wrapper.h

@@ -17,6 +17,10 @@
 #ifndef __ONERT_API_PYTHON_NNFW_API_WRAPPER_H__
 #define __ONERT_API_PYTHON_NNFW_API_WRAPPER_H__
 
+#include <stdexcept>
+#include <string>
+#include <vector>
+
 #include "nnfw.h"
 #include "nnfw_experimental.h"
 #include "nnfw_internal.h"
@@ -34,28 +38,58 @@ namespace python
 namespace py = pybind11;
 
 /**
- *  @brief  tensor info describes the type and shape of tensors
+ * @brief Data type used by OneRT API.
+ *
+ * This structure maps NNFW_TYPE and numpy dtype.
+ */
+struct dtype
+{
+  NNFW_TYPE type;
+  py::dtype dtype;
+  // The name of the dtype, e.g., "float32", "int32", etc.
+  // This is mainly for __repr__ implementation.
+  const char *name;
+};
+
+/**
+ *  @brief This class describes the type and shape of tensors.
  *
  * This structure is used to describe input and output tensors.
  * Application can get input and output tensor type and shape described in model by using
  * {@link input_tensorinfo} and {@link output_tensorinfo}
  *
- * Maximum rank is 6 (NNFW_MAX_RANK).
- * And tensor's dimension value is filled in 'dims' field from index 0.
- * For example, if tensor's rank is 4,
- * application can get dimension value from dims[0], dims[1], dims[2], and dims[3]
+ * Tensor's dimension values are stored in the `_shape` vector. The size of the
+ * vector determines the rank of the tensor. The maximum rank is 6 (NNFW_MAX_RANK).
+ *
+ * This class is immutable. To change the tensorinfo, create a new instance
+ * and use {@link set_input_tensorinfo} to set it.
  */
-struct tensorinfo
+class tensorinfo
 {
-  /** The data type */
-  const char *dtype;
-  /** The number of dimensions (rank) */
-  int32_t rank;
+public:
+  using SHAPE = std::vector<int32_t>;
+
+private:
+  /** The data type. */
+  dtype _dtype;
   /**
-   * The dimension of tensor.
-   * Maximum rank is 6 (NNFW_MAX_RANK).
+   * Shape of the tensor.
+   * The maximum allowed rank is 6 (NNFW_MAX_RANK).
    */
-  int32_t dims[NNFW_MAX_RANK];
+  SHAPE _shape;
+
+public:
+  tensorinfo(dtype dt, SHAPE shape) : _dtype(dt)
+  {
+    if (shape.size() > NNFW_MAX_RANK)
+      throw std::invalid_argument(std::string("Rank of tensor must be less than or equal to ") +
+                                  std::to_string(NNFW_MAX_RANK));
+    _shape = std::move(shape);
+  }
+
+  const auto &get_dtype() const { return _dtype; }
+  const auto &get_shape() const { return _shape; }
+  size_t get_rank() const { return _shape.size(); }
 };
 
 /**
@@ -76,20 +110,12 @@ void ensure_status(NNFW_STATUS status);
 NNFW_LAYOUT getLayout(const char *layout = "");
 
 /**
- * Convert the type with string to NNFW_TYPE
+ * Convert the NNFW_TYPE type to the dtype type.
  *
- * @param[in] type type to be converted
- * @return proper type if exists
+ * @param[in] type The type to be converted.
+ * @return Corresponding dtype type.
  */
-NNFW_TYPE getType(const char *type = "");
-
-/**
- * Convert the type with NNFW_TYPE to string
- *
- * @param[in] type type to be converted
- * @return proper type
- */
-const char *getStringType(NNFW_TYPE type);
+dtype get_dtype(NNFW_TYPE type);
 
 /**
  * @brief     Get the total number of elements in nnfw_tensorinfo->dims.
@@ -101,26 +127,6 @@ const char *getStringType(NNFW_TYPE type);
  */
 uint64_t num_elems(const nnfw_tensorinfo *tensor_info);
 
-/**
- * @brief     Get nnfw_tensorinfo->dims.
- *
- * This function is called to get dimension array of tensorinfo.
- *
- * @param[in] tensor_info Tensor info (shape, type, etc)
- * @return python list of dims
- */
-py::list get_dims(const tensorinfo &tensor_info);
-
-/**
- * @brief     Set nnfw_tensorinfo->dims.
- *
- * This function is called to set dimension array of tensorinfo.
- *
- * @param[in] tensor_info Tensor info (shape, type, etc)
- * @param[in] array       array to set dimension
- */
-void set_dims(tensorinfo &tensor_info, const py::list &array);
-
 class NNFW_SESSION
 {
 private:

runtime/onert/api/python/package/__init__.py

@@ -1,11 +1,12 @@
 # Define the public API of the onert package
-__all__ = ["infer", "tensorinfo", "experimental"]
+__all__ = ["dtype", "infer", "tensorinfo", "experimental"]
+
+# Import and expose tensorinfo and tensor data types
+from .native.libnnfw_api_pybind import dtype, tensorinfo
+from .native.libnnfw_api_pybind.dtypes import *
 
 # Import and expose the infer module's functionalities
 from . import infer
 
-# Import and expose tensorinfo
-from .common import tensorinfo as tensorinfo
-
 # Import and expose the experimental module's functionalities
 from . import experimental

runtime/onert/api/python/package/common/basesession.py

@@ -1,15 +1,15 @@
 from typing import List
 import numpy as np
 
-from ..native.libnnfw_api_pybind import infer, tensorinfo
+from ..native.libnnfw_api_pybind import tensorinfo
 from ..native.libnnfw_api_pybind.exception import OnertError
 
 
 def num_elems(tensor_info):
-    """Get the total number of elements in nnfw_tensorinfo.dims."""
+    """Get the total number of elements in tensorinfo.shape."""
     n = 1
-    for x in range(tensor_info.rank):
-        n *= tensor_info.dims[x]
+    for x in tensor_info.shape:
+        n *= x
     return n
 
 
@@ -131,13 +131,12 @@ def set_inputs(self, size, inputs_array=[]):
                 input_array = np.zeros((num_elems(input_tensorinfo)),
                                        dtype=input_tensorinfo.dtype)
 
-            # Check if the shape of input_array matches the dims of input_tensorinfo
-            if input_array.shape != tuple(input_tensorinfo.dims):
+            # Check if the shape of input_array matches the input_tensorinfo
+            if input_array.shape != input_tensorinfo.shape:
                 # If not, set the input tensor info to match the input_array shape
                 try:
-                    input_tensorinfo.rank = len(input_array)
-                    input_tensorinfo.dims = list(input_array.shape)
-                    self.session.set_input_tensorinfo(i, input_tensorinfo)
+                    info = tensorinfo(input_tensorinfo.dtype, input_array)
+                    self.session.set_input_tensorinfo(i, info)
                 except Exception as e:
                     raise OnertError(f"Failed to set input tensor info #{i}: {e}") from e
 
@@ -176,18 +175,3 @@ def _set_outputs(self, size):
                 raise OnertError(f"Failed to get output #{i}: {e}") from e
 
             self.outputs.append(output_array)
-
-
-def tensorinfo():
-    """
-    Shortcut to create a fresh tensorinfo instance.
-    Raises:
-        OnertError: If the C-API call fails.
-    """
-
-    try:
-        return infer.nnfw_tensorinfo()
-    except OnertError:
-        raise
-    except Exception as e:
-        raise OnertError(f"Failed to create tensorinfo: {e}") from e

runtime/onert/api/python/package/infer/session.py

@@ -68,27 +68,26 @@ def infer(
                     fixed_infos = []
                     for idx, info in enumerate(original_infos):
                         input_shape = inputs_array[idx].shape
-                        new_dims = []
+                        new_shape = []
                         static_dim_changed = False
                         # only the first `info.rank` entries matter
-                        for j, d in enumerate(info.dims[:info.rank]):
+                        for j, d in enumerate(info.shape):
                             if d == -1:
                                 # replace dynamic dim with actual incoming shape
-                                new_dims.append(input_shape[j])
+                                new_shape.append(input_shape[j])
                             elif d == input_shape[j]:
                                 # static dim must match the provided array
-                                new_dims.append(d)
+                                new_shape.append(d)
                             else:
                                 static_dim_changed = True
 
                         if static_dim_changed:
                             warnings.warn(
                                 f"infer() called with input {idx}'s shape={input_shape}, "
-                                f"which differs from model's expected shape={tuple(info.dims)}. "
+                                f"which differs from model's expected shape={info.shape}. "
                                 "Ensure this is intended.", UserWarning)
 
-                        info.dims = new_dims
-                        fixed_infos.append(info)
+                        fixed_infos.append(tensorinfo(info.dtype, new_shape))
 
                     # Update tensorinfo to optimize using it
                     self._update_inputs_tensorinfo(fixed_infos)
@@ -151,10 +150,10 @@ def _update_inputs_tensorinfo(self, new_infos: List[tensorinfo]) -> None:
 
         for i, info in enumerate(new_infos):
             # Check for any negative dimension in the specified rank
-            if any(d < 0 for d in info.dims[:info.rank]):
+            if any(d < 0 for d in info.shape):
                 raise ValueError(
                     f"Input tensorinfo at index {i} contains negative dimension(s): "
-                    f"{info.dims[:info.rank]}")
+                    f"{info.shape}")
             try:
                 self.session.set_input_tensorinfo(i, info)
             except ValueError:

runtime/onert/api/python/src/bindings/nnfw_tensorinfo_bindings.cc

@@ -16,25 +16,63 @@
 
 #include "nnfw_tensorinfo_bindings.h"
 
+#include <cstdint>
+#include <vector>
+
 #include "nnfw_api_wrapper.h"
 
 namespace onert::api::python
 {
 
 namespace py = pybind11;
 
-// Bind the `tensorinfo` class
+template <typename T> py::tuple shape_to_tuple(const T &shape)
+{
+  py::tuple t(shape.size());
+  for (size_t i = 0; i < shape.size(); i++)
+    t[i] = shape[i];
+  return t;
+}
+
+// Bind the `tensorinfo` class and related `dtype` class.
 void bind_tensorinfo(py::module_ &m)
 {
-  py::class_<tensorinfo>(m, "tensorinfo", "tensorinfo describes the type and shape of tensors",
+  py::class_<dtype>(m, "dtype", "Defines the type of the OneRT tensor.", py::module_local())
+    .def("__repr__", [](const dtype &dt) { return std::string("onert.") + dt.name; })
+    .def_readonly("dtype", &dtype::dtype, "A numpy data type.");
+
+  py::class_<tensorinfo>(m, "tensorinfo",
+                         "Immutable information about the type and shape of a tensor.",
                          py::module_local())
-    .def(py::init<>(), "The constructor of tensorinfo")
-    .def_readwrite("dtype", &tensorinfo::dtype, "The data type")
-    .def_readwrite("rank", &tensorinfo::rank, "The number of dimensions (rank)")
-    .def_property(
-      "dims", [](const tensorinfo &ti) { return get_dims(ti); },
-      [](tensorinfo &ti, const py::list &dims_list) { set_dims(ti, dims_list); },
-      "The dimension of tensor. Maximum rank is 6 (NNFW_MAX_RANK).");
+    .def(py::init<struct dtype, tensorinfo::SHAPE>(),
+         "Initialize new tensorinfo with dtype and shape.", py::arg("dtype"), py::arg("shape"))
+    .def_property_readonly("dtype", &tensorinfo::get_dtype, "The data type of the tensor.")
+    .def_property_readonly("rank", &tensorinfo::get_rank,
+                           "The rank of the tensor. The maximum supported rank is 6.")
+    .def_property_readonly(
+      "shape", [](const tensorinfo &ti) { return shape_to_tuple(ti.get_shape()); },
+      "The shape of the tensor.")
+    .def("__repr__", [](const tensorinfo &ti) {
+      auto dtype = py::repr(py::cast(ti.get_dtype())).cast<std::string>();
+      auto shape = py::repr(shape_to_tuple(ti.get_shape())).cast<std::string>();
+      return "<tensorinfo dtype=" + dtype + " shape=" + shape + ">";
+    });
+
+  static const dtype dtypes[] = {
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_FLOAT32),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_INT32),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_QUANT8_ASYMM),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_UINT8),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_BOOL),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_INT64),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_QUANT8_ASYMM_SIGNED),
+    get_dtype(NNFW_TYPE::NNFW_TYPE_TENSOR_QUANT16_SYMM_SIGNED),
+  };
+
+  // Export OneRT dtypes in a submodule
+  auto m_dtypes = m.def_submodule("dtypes", "OneRT tensor data types");
+  for (const auto &dt : dtypes)
+    m_dtypes.attr(dt.name) = dt;
 }
 
 void bind_nnfw_enums(py::module_ &m)