Expose DType size through the C api (#39)

* Expose dtype size * rebase + move to new MLX C API --------- Co-authored-by: Ronan Collobert <[email protected]>
ml-explore · Nov 19, 2024 · 43c06a7 · 43c06a7
1 parent 7999777
commit 43c06a7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/mlx/c/array.cpp b/mlx/c/array.cpp
@@ -7,6 +7,10 @@
 #include "mlx/c/private/mlx.h"
 #include "mlx/c/string.h"
 
+extern "C" size_t mlx_dtype_size(mlx_dtype dtype) {
+  return mlx_dtype_to_cpp(dtype).size();
+}
+
 extern "C" int mlx_array_tostring(mlx_string* str_, const mlx_array arr) {
   try {
     std::ostringstream os;

diff --git a/mlx/c/array.h b/mlx/c/array.h
@@ -49,6 +49,11 @@ typedef enum mlx_dtype_ {
   MLX_COMPLEX64,
 } mlx_dtype;
 
+/**
+ * Size of given mlx_dtype datatype in bytes.
+ */
+size_t mlx_dtype_size(mlx_dtype dtype);
+
 /**
  * Get array description.
  */