add datatype enum to python module (#404)

* add datatype enum to python module Signed-off-by: koubaa <[email protected]> * add datatype enum to python module Signed-off-by: koubaa <[email protected]> * add doc Signed-off-by: koubaa <[email protected]> --------- Signed-off-by: koubaa <[email protected]> Co-authored-by: koubaa <[email protected]>
KomputeProject · Dec 10, 2024 · 48c127d · 48c127d
1 parent cb85d21
commit 48c127d
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -187,6 +187,7 @@ def kompute(shader):
     # Explicit type constructor supports uint32, int32, double, float and bool
     tensor_out_a = mgr.tensor_t(np.array([0, 0, 0], dtype=np.uint32))
     tensor_out_b = mgr.tensor_t(np.array([0, 0, 0], dtype=np.uint32))
+    assert(t_data.data_type() == kp.DataTypes.uint)
 
     params = [tensor_in_a, tensor_in_b, tensor_out_a, tensor_out_b]
 

diff --git a/docs/overview/python-examples.rst b/docs/overview/python-examples.rst
@@ -63,7 +63,10 @@ Similarly you can find the same extended example as above:
     # Can be initialized with List[] or np.Array
     tensor_in_a = mgr.tensor([2, 2, 2])
     tensor_in_b = mgr.tensor([1, 2, 3])
-    tensor_out = mgr.tensor([0, 0, 0])
+
+    # By default, tensors use a float type, but that can be explicitly specified
+    tensor_out = mgr.tensor_t([0, 0, 0], dtype=np.float32)
+    assert(tensor_out.data_type() == kp.DataTypes.float)
 
     seq = mgr.sequence()
     seq.eval(kp.OpTensorSyncDevice([tensor_in_a, tensor_in_b, tensor_out]))

diff --git a/docs/overview/python-reference.rst b/docs/overview/python-reference.rst
@@ -39,3 +39,16 @@ TensorType
 .. automodule:: kp
    :members:
 
+MemoryTypes
+-------
+
+.. automodule:: kp
+   :members:
+
+DataTypes
+-------
+
+.. automodule:: kp
+   :members:
+
+
diff --git a/python/src/main.cpp b/python/src/main.cpp
@@ -65,6 +65,39 @@ PYBIND11_MODULE(kp, m)
 
     py::module_ np = py::module_::import("numpy");
 
+    py::enum_<kp::Memory::DataTypes>(m, "DataTypes")
+      .value("bool",
+             kp::Memory::DataTypes::eBool,
+             DOC(kp, Memory, DataTypes, eBool))
+      .value("int",
+             kp::Memory::DataTypes::eInt,
+             DOC(kp, Memory, DataTypes, eInt))
+      .value("uint",
+             kp::Memory::DataTypes::eUnsignedInt,
+             DOC(kp, Memory, DataTypes, eUnsignedInt))
+      .value("float",
+             kp::Memory::DataTypes::eFloat,
+             DOC(kp, Memory, DataTypes, eFloat))
+      .value("double",
+             kp::Memory::DataTypes::eDouble,
+             DOC(kp, Memory, DataTypes, eDouble))
+      .value("custom",
+             kp::Memory::DataTypes::eCustom,
+             DOC(kp, Memory, DataTypes, eCustom))
+      .value("short",
+             kp::Memory::DataTypes::eShort,
+             DOC(kp, Memory, DataTypes, eShort))
+      .value("ushort",
+             kp::Memory::DataTypes::eUnsignedShort,
+             DOC(kp, Memory, DataTypes, eUnsignedShort))
+      .value("char",
+             kp::Memory::DataTypes::eChar,
+             DOC(kp, Memory, DataTypes, eChar))
+      .value("uchar",
+             kp::Memory::DataTypes::eUnsignedChar,
+             DOC(kp, Memory, DataTypes, eUnsignedChar))
+      .export_values();
+
     py::enum_<kp::Memory::MemoryTypes>(m, "MemoryTypes")
       .value("device",
              kp::Memory::MemoryTypes::eDevice,