| 98 | }; |
| 99 | |
| 100 | void init_array(nb::module_& m) { |
| 101 | // Types |
| 102 | nb::class_<mx::Dtype>( |
| 103 | m, |
| 104 | "Dtype", |
| 105 | R"pbdoc( |
| 106 | An object to hold the type of a :class:`array`. |
| 107 | |
| 108 | See the :ref:`list of types <data_types>` for more details |
| 109 | on available data types. |
| 110 | )pbdoc") |
| 111 | .def_prop_ro( |
| 112 | "size", &mx::Dtype::size, R"pbdoc(Size of the type in bytes.)pbdoc") |
| 113 | .def( |
| 114 | "__repr__", |
| 115 | [](const mx::Dtype& t) { |
| 116 | std::ostringstream os; |
| 117 | os << "mlx.core."; |
| 118 | os << t; |
| 119 | return os.str(); |
| 120 | }) |
| 121 | .def( |
| 122 | "__eq__", |
| 123 | [](const mx::Dtype& t, const nb::object& other) { |
| 124 | return nb::isinstance<mx::Dtype>(other) && |
| 125 | t == nb::cast<mx::Dtype>(other); |
| 126 | }) |
| 127 | .def("__hash__", [](const mx::Dtype& t) { |
| 128 | return static_cast<int64_t>(t.val()); |
| 129 | }); |
| 130 | |
| 131 | m.attr("bool_") = nb::cast(mx::bool_); |
| 132 | m.attr("uint8") = nb::cast(mx::uint8); |
| 133 | m.attr("uint16") = nb::cast(mx::uint16); |
| 134 | m.attr("uint32") = nb::cast(mx::uint32); |
| 135 | m.attr("uint64") = nb::cast(mx::uint64); |
| 136 | m.attr("int8") = nb::cast(mx::int8); |
| 137 | m.attr("int16") = nb::cast(mx::int16); |
| 138 | m.attr("int32") = nb::cast(mx::int32); |
| 139 | m.attr("int64") = nb::cast(mx::int64); |
| 140 | m.attr("float16") = nb::cast(mx::float16); |
| 141 | m.attr("float32") = nb::cast(mx::float32); |
| 142 | m.attr("float64") = nb::cast(mx::float64); |
| 143 | m.attr("bfloat16") = nb::cast(mx::bfloat16); |
| 144 | m.attr("complex64") = nb::cast(mx::complex64); |
| 145 | nb::enum_<mx::Dtype::Category>( |
| 146 | m, |
| 147 | "DtypeCategory", |
| 148 | R"pbdoc( |
| 149 | Type to hold categories of :class:`dtypes <Dtype>`. |
| 150 | |
| 151 | * :attr:`~mlx.core.generic` |
| 152 | |
| 153 | * :ref:`bool_ <data_types>` |
| 154 | * :attr:`~mlx.core.number` |
| 155 | |
| 156 | * :attr:`~mlx.core.integer` |
| 157 |
no test coverage detected