| |
| #include <cstdint> |
| #include <cstring> |
| #include <sstream> |
|
|
| #include <nanobind/ndarray.h> |
| #include <nanobind/stl/complex.h> |
| #include <nanobind/stl/optional.h> |
| #include <nanobind/stl/string.h> |
| #include <nanobind/stl/variant.h> |
| #include <nanobind/stl/vector.h> |
| #include <nanobind/typing.h> |
|
|
| #include "mlx/backend/metal/metal.h" |
| #include "python/src/buffer.h" |
| #include "python/src/convert.h" |
| #include "python/src/indexing.h" |
| #include "python/src/small_vector.h" |
| #include "python/src/utils.h" |
|
|
| #include "mlx/mlx.h" |
|
|
| namespace mx = mlx::core; |
| namespace nb = nanobind; |
| using namespace nb::literals; |
|
|
| class ArrayAt { |
| public: |
| ArrayAt(mx::array x) : x_(std::move(x)) {} |
| ArrayAt& set_indices(nb::object indices) { |
| initialized_ = true; |
| indices_ = indices; |
| return *this; |
| } |
| void check_initialized() { |
| if (!initialized_) { |
| throw std::invalid_argument( |
| "Must give indices to array.at (e.g. `x.at[0].add(4)`)."); |
| } |
| } |
|
|
| mx::array add(const ScalarOrArray& v) { |
| check_initialized(); |
| return mlx_add_item(x_, indices_, v); |
| } |
| mx::array subtract(const ScalarOrArray& v) { |
| check_initialized(); |
| return mlx_subtract_item(x_, indices_, v); |
| } |
| mx::array multiply(const ScalarOrArray& v) { |
| check_initialized(); |
| return mlx_multiply_item(x_, indices_, v); |
| } |
| mx::array divide(const ScalarOrArray& v) { |
| check_initialized(); |
| return mlx_divide_item(x_, indices_, v); |
| } |
| mx::array maximum(const ScalarOrArray& v) { |
| check_initialized(); |
| return mlx_maximum_item(x_, indices_, v); |
| } |
| mx::array minimum(const ScalarOrArray& v) { |
| check_initialized(); |
| return mlx_minimum_item(x_, indices_, v); |
| } |
|
|
| private: |
| mx::array x_; |
| bool initialized_{false}; |
| nb::object indices_; |
| }; |
|
|
| class ArrayPythonIterator { |
| public: |
| ArrayPythonIterator(mx::array x) : idx_(0), x_(std::move(x)) { |
| if (x_.shape(0) > 0 && x_.shape(0) < 10) { |
| splits_ = mx::split(x_, x_.shape(0)); |
| } |
| } |
|
|
| mx::array next() { |
| if (idx_ >= x_.shape(0)) { |
| throw nb::stop_iteration(); |
| } |
|
|
| if (idx_ >= 0 && idx_ < splits_.size()) { |
| return mx::squeeze(splits_[idx_++], 0); |
| } |
|
|
| return *(x_.begin() + idx_++); |
| } |
|
|
| private: |
| int idx_; |
| mx::array x_; |
| std::vector<mx::array> splits_; |
| }; |
|
|
| void init_array(nb::module_& m) { |
| |
| mx::get_global_formatter().capitalize_bool = true; |
|
|
| |
| nb::class_<mx::Dtype>( |
| m, |
| "Dtype", |
| R"pbdoc( |
| An object to hold the type of a :class:`array`. |
| |
| See the :ref:`list of types <data_types>` for more details |
| on available data types. |
| )pbdoc") |
| .def_prop_ro( |
| "size", &mx::Dtype::size, R"pbdoc(Size of the type in bytes.)pbdoc") |
| .def( |
| "__repr__", |
| [](const mx::Dtype& t) { |
| std::ostringstream os; |
| os << "mlx.core."; |
| os << t; |
| return os.str(); |
| }) |
| .def( |
| "__eq__", |
| [](const mx::Dtype& t, const nb::object& other) { |
| return nb::isinstance<mx::Dtype>(other) && |
| t == nb::cast<mx::Dtype>(other); |
| }) |
| .def("__hash__", [](const mx::Dtype& t) { |
| return static_cast<int64_t>(t.val()); |
| }); |
|
|
| m.attr("bool_") = nb::cast(mx::bool_); |
| m.attr("uint8") = nb::cast(mx::uint8); |
| m.attr("uint16") = nb::cast(mx::uint16); |
| m.attr("uint32") = nb::cast(mx::uint32); |
| m.attr("uint64") = nb::cast(mx::uint64); |
| m.attr("int8") = nb::cast(mx::int8); |
| m.attr("int16") = nb::cast(mx::int16); |
| m.attr("int32") = nb::cast(mx::int32); |
| m.attr("int64") = nb::cast(mx::int64); |
| m.attr("float16") = nb::cast(mx::float16); |
| m.attr("float32") = nb::cast(mx::float32); |
| m.attr("float64") = nb::cast(mx::float64); |
| m.attr("bfloat16") = nb::cast(mx::bfloat16); |
| m.attr("complex64") = nb::cast(mx::complex64); |
| nb::enum_<mx::Dtype::Category>( |
| m, |
| "DtypeCategory", |
| R"pbdoc( |
| Type to hold categories of :class:`dtypes <Dtype>`. |
| |
| * :attr:`~mlx.core.generic` |
| |
| * :ref:`bool_ <data_types>` |
| * :attr:`~mlx.core.number` |
| |
| * :attr:`~mlx.core.integer` |
| |
| * :attr:`~mlx.core.unsignedinteger` |
| |
| * :ref:`uint8 <data_types>` |
| * :ref:`uint16 <data_types>` |
| * :ref:`uint32 <data_types>` |
| * :ref:`uint64 <data_types>` |
| |
| * :attr:`~mlx.core.signedinteger` |
| |
| * :ref:`int8 <data_types>` |
| * :ref:`int32 <data_types>` |
| * :ref:`int64 <data_types>` |
| |
| * :attr:`~mlx.core.inexact` |
| |
| * :attr:`~mlx.core.floating` |
| |
| * :ref:`float16 <data_types>` |
| * :ref:`bfloat16 <data_types>` |
| * :ref:`float32 <data_types>` |
| * :ref:`float64 <data_types>` |
| |
| * :attr:`~mlx.core.complexfloating` |
| |
| * :ref:`complex64 <data_types>` |
| |
| See also :func:`~mlx.core.issubdtype`. |
| )pbdoc") |
| .value("complexfloating", mx::complexfloating) |
| .value("floating", mx::floating) |
| .value("inexact", mx::inexact) |
| .value("signedinteger", mx::signedinteger) |
| .value("unsignedinteger", mx::unsignedinteger) |
| .value("integer", mx::integer) |
| .value("number", mx::number) |
| .value("generic", mx::generic) |
| .export_values(); |
|
|
| nb::class_<mx::finfo>( |
| m, |
| "finfo", |
| R"pbdoc( |
| Get information on floating-point types. |
| )pbdoc") |
| .def(nb::init<mx::Dtype>()) |
| .def_ro( |
| "min", |
| &mx::finfo::min, |
| R"pbdoc(The smallest representable number.)pbdoc") |
| .def_ro( |
| "max", |
| &mx::finfo::max, |
| R"pbdoc(The largest representable number.)pbdoc") |
| .def_ro( |
| "eps", |
| &mx::finfo::eps, |
| R"pbdoc( |
| The difference between 1.0 and the next smallest |
| representable number larger than 1.0. |
| )pbdoc") |
| .def_ro("dtype", &mx::finfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc") |
| .def("__repr__", [](const mx::finfo& f) { |
| std::ostringstream os; |
| os << "finfo(" |
| << "min=" << f.min << ", max=" << f.max << ", dtype=" << f.dtype |
| << ")"; |
| return os.str(); |
| }); |
|
|
| nb::class_<mx::iinfo>( |
| m, |
| "iinfo", |
| R"pbdoc( |
| Get information on integer types. |
| )pbdoc") |
| .def(nb::init<mx::Dtype>()) |
| .def_ro( |
| "min", |
| &mx::iinfo::min, |
| R"pbdoc(The smallest representable number.)pbdoc") |
| .def_ro( |
| "max", |
| &mx::iinfo::max, |
| R"pbdoc(The largest representable number.)pbdoc") |
| .def_ro("dtype", &mx::iinfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc") |
| .def("__repr__", [](const mx::iinfo& i) { |
| std::ostringstream os; |
| os << "iinfo(" |
| << "min=" << i.min << ", max=" << i.max << ", dtype=" << i.dtype |
| << ")"; |
| return os.str(); |
| }); |
|
|
| nb::class_<ArrayAt>( |
| m, |
| "ArrayAt", |
| R"pbdoc( |
| A helper object to apply updates at specific indices. |
| )pbdoc") |
| .def("__getitem__", &ArrayAt::set_indices, "indices"_a.none()) |
| .def("add", &ArrayAt::add, "value"_a) |
| .def("subtract", &ArrayAt::subtract, "value"_a) |
| .def("multiply", &ArrayAt::multiply, "value"_a) |
| .def("divide", &ArrayAt::divide, "value"_a) |
| .def("maximum", &ArrayAt::maximum, "value"_a) |
| .def("minimum", &ArrayAt::minimum, "value"_a); |
|
|
| nb::class_<ArrayLike>( |
| m, |
| "ArrayLike", |
| R"pbdoc( |
| Any Python object which has an ``__mlx__array__`` method that |
| returns an :obj:`array`. |
| )pbdoc") |
| .def(nb::init_implicit<nb::object>()); |
|
|
| nb::class_<ArrayPythonIterator>( |
| m, |
| "ArrayIterator", |
| R"pbdoc( |
| A helper object to iterate over the 1st dimension of an array. |
| )pbdoc") |
| .def("__next__", &ArrayPythonIterator::next) |
| .def("__iter__", [](const ArrayPythonIterator& it) { return it; }); |
|
|
| |
| PyType_Slot array_slots[] = { |
| {Py_bf_getbuffer, (void*)getbuffer}, |
| {Py_bf_releasebuffer, (void*)releasebuffer}, |
| {0, nullptr}}; |
|
|
| nb::class_<mx::array>( |
| m, |
| "array", |
| R"pbdoc(An N-dimensional array object.)pbdoc", |
| nb::type_slots(array_slots), |
| nb::is_weak_referenceable()) |
| .def( |
| "__init__", |
| [](mx::array* aptr, ArrayInitType v, std::optional<mx::Dtype> t) { |
| new (aptr) mx::array(create_array(v, t)); |
| }, |
| "val"_a, |
| "dtype"_a = nb::none(), |
| nb::sig( |
| "def __init__(self: array, val: Union[scalar, list, tuple, numpy.ndarray, array], dtype: Optional[Dtype] = None)")) |
| .def_prop_ro( |
| "size", |
| &mx::array::size, |
| R"pbdoc(Number of elements in the array.)pbdoc") |
| .def_prop_ro( |
| "ndim", &mx::array::ndim, R"pbdoc(The array's dimension.)pbdoc") |
| .def_prop_ro( |
| "itemsize", |
| &mx::array::itemsize, |
| R"pbdoc(The size of the array's datatype in bytes.)pbdoc") |
| .def_prop_ro( |
| "nbytes", |
| &mx::array::nbytes, |
| R"pbdoc(The number of bytes in the array.)pbdoc") |
| .def_prop_ro( |
| "shape", |
| [](const mx::array& a) { return nb::cast(a.shape()); }, |
| nb::sig("def shape(self) -> tuple[int, ...]"), |
| R"pbdoc( |
| The shape of the array as a Python tuple. |
| |
| Returns: |
| tuple(int): A tuple containing the sizes of each dimension. |
| )pbdoc") |
| .def_prop_ro( |
| "dtype", |
| &mx::array::dtype, |
| R"pbdoc( |
| The array's :class:`Dtype`. |
| )pbdoc") |
| .def_prop_ro( |
| "real", |
| [](const mx::array& a) { return mx::real(a); }, |
| R"pbdoc( |
| The real part of a complex array. |
| )pbdoc") |
| .def_prop_ro( |
| "imag", |
| [](const mx::array& a) { return mx::imag(a); }, |
| R"pbdoc( |
| The imaginary part of a complex array. |
| )pbdoc") |
| .def( |
| "item", |
| &to_scalar, |
| nb::sig("def item(self) -> scalar"), |
| R"pbdoc( |
| Access the value of a scalar array. |
| |
| Returns: |
| Standard Python scalar. |
| )pbdoc") |
| .def( |
| "tolist", |
| &tolist, |
| nb::sig("def tolist(self) -> list_or_scalar"), |
| R"pbdoc( |
| Convert the array to a Python :class:`list`. |
| |
| Returns: |
| list: The Python list. |
| |
| If the array is a scalar then a standard Python scalar is returned. |
| |
| If the array has more than one dimension then the result is a nested |
| list of lists. |
| |
| The value type of the list corresponding to the last dimension is either |
| ``bool``, ``int`` or ``float`` depending on the ``dtype`` of the array. |
| )pbdoc") |
| .def( |
| "astype", |
| &mx::astype, |
| "dtype"_a, |
| "stream"_a = nb::none(), |
| R"pbdoc( |
| Cast the array to a specified type. |
| |
| Args: |
| dtype (Dtype): Type to which the array is cast. |
| stream (Stream): Stream (or device) for the operation. |
| |
| Returns: |
| array: The array with type ``dtype``. |
| )pbdoc") |
| .def( |
| "__array_namespace__", |
| [](const mx::array& a, |
| const std::optional<std::string>& api_version) { |
| if (api_version) { |
| throw std::invalid_argument( |
| "Explicitly specifying api_version is not yet implemented."); |
| } |
| return nb::module_::import_("mlx.core"); |
| }, |
| "api_version"_a = nb::none(), |
| R"pbdoc( |
| Returns an object that has all the array API functions on it. |
| |
| See the `Python array API <https://data-apis.org/array-api/latest/index.html>`_ |
| for more information. |
| |
| Args: |
| api_version (str, optional): String representing the version |
| of the array API spec to return. Default: ``None``. |
| |
| Returns: |
| out (Any): An object representing the array API namespace. |
| )pbdoc") |
| .def("__getitem__", mlx_get_item, nb::arg().none()) |
| .def("__setitem__", mlx_set_item, nb::arg().none(), nb::arg()) |
| .def_prop_ro( |
| "at", |
| [](const mx::array& a) { return ArrayAt(a); }, |
| R"pbdoc( |
| Used to apply updates at the given indices. |
| |
| .. note:: |
| |
| Regular in-place updates map to assignment. For instance ``x[idx] += y`` |
| maps to ``x[idx] = x[idx] + y``. As a result, assigning to the |
| same index ignores all but one update. Using ``x.at[idx].add(y)`` |
| will correctly apply all updates to all indices. |
| |
| .. list-table:: |
| :header-rows: 1 |
| |
| * - array.at syntax |
| - In-place syntax |
| * - ``x = x.at[idx].add(y)`` |
| - ``x[idx] += y`` |
| * - ``x = x.at[idx].subtract(y)`` |
| - ``x[idx] -= y`` |
| * - ``x = x.at[idx].multiply(y)`` |
| - ``x[idx] *= y`` |
| * - ``x = x.at[idx].divide(y)`` |
| - ``x[idx] /= y`` |
| * - ``x = x.at[idx].maximum(y)`` |
| - ``x[idx] = mx.maximum(x[idx], y)`` |
| * - ``x = x.at[idx].minimum(y)`` |
| - ``x[idx] = mx.minimum(x[idx], y)`` |
| |
| Example: |
| >>> a = mx.array([0, 0]) |
| >>> idx = mx.array([0, 1, 0, 1]) |
| >>> a[idx] += 1 |
| >>> a |
| array([1, 1], dtype=int32) |
| >>> |
| >>> a = mx.array([0, 0]) |
| >>> a.at[idx].add(1) |
| array([2, 2], dtype=int32) |
| )pbdoc") |
| .def( |
| "__len__", |
| [](const mx::array& a) { |
| if (a.ndim() == 0) { |
| throw nb::type_error("len() 0-dimensional array."); |
| } |
| return a.shape(0); |
| }) |
| .def( |
| "__iter__", [](const mx::array& a) { return ArrayPythonIterator(a); }) |
| .def( |
| "__getstate__", |
| [](const mx::array& a) { |
| auto nd = (a.dtype() == mx::bfloat16) |
| ? mlx_to_np_array(mx::view(a, mx::uint16)) |
| : mlx_to_np_array(a); |
| return nb::make_tuple(nd, static_cast<uint8_t>(a.dtype().val())); |
| }) |
| .def( |
| "__setstate__", |
| [](mx::array& arr, const nb::tuple& state) { |
| if (nb::len(state) != 2) { |
| throw std::invalid_argument( |
| "Invalid pickle state: expected (ndarray, Dtype::Val)"); |
| } |
| using ND = nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>; |
| ND nd = nb::cast<ND>(state[0]); |
| auto val = static_cast<mx::Dtype::Val>(nb::cast<uint8_t>(state[1])); |
| if (val == mx::Dtype::Val::bfloat16) { |
| auto owner = nb::handle(state[0].ptr()); |
| new (&arr) mx::array(nd_array_to_mlx( |
| ND(nd.data(), |
| nd.ndim(), |
| reinterpret_cast<const size_t*>(nd.shape_ptr()), |
| owner, |
| nullptr, |
| nb::bfloat16), |
| mx::bfloat16)); |
| } else { |
| new (&arr) mx::array(nd_array_to_mlx(nd, std::nullopt)); |
| } |
| }) |
| .def("__dlpack__", [](const mx::array& a) { return mlx_to_dlpack(a); }) |
| .def( |
| "__dlpack_device__", |
| [](const mx::array& a) { |
| |
| |
| if (mx::metal::is_available()) { |
| return nb::make_tuple(8, 0); |
| } else if (mx::cu::is_available()) { |
| return nb::make_tuple(13, 0); |
| } else { |
| |
| return nb::make_tuple(1, 0); |
| } |
| }) |
| .def("__copy__", [](const mx::array& self) { return mx::array(self); }) |
| .def( |
| "__deepcopy__", |
| [](const mx::array& self, nb::dict) { return mx::array(self); }, |
| "memo"_a) |
| .def( |
| "__add__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("addition", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| return mx::add(a, b); |
| }, |
| "other"_a) |
| .def( |
| "__iadd__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace addition", v); |
| } |
| a.overwrite_descriptor(mx::add(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__radd__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("addition", v); |
| } |
| return mx::add(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__sub__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("subtraction", v); |
| } |
| return mx::subtract(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__isub__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace subtraction", v); |
| } |
| a.overwrite_descriptor(mx::subtract(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__rsub__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("subtraction", v); |
| } |
| return mx::subtract(to_array(v, a.dtype()), a); |
| }, |
| "other"_a) |
| .def( |
| "__mul__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("multiplication", v); |
| } |
| return mx::multiply(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__imul__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace multiplication", v); |
| } |
| a.overwrite_descriptor(mx::multiply(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__rmul__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("multiplication", v); |
| } |
| return mx::multiply(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__truediv__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("division", v); |
| } |
| return mx::divide(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__itruediv__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace division", v); |
| } |
| if (!mx::issubdtype(a.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "In place division cannot cast to non-floating point type."); |
| } |
| a.overwrite_descriptor(divide(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__rtruediv__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("division", v); |
| } |
| return mx::divide(to_array(v, a.dtype()), a); |
| }, |
| "other"_a) |
| .def( |
| "__div__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("division", v); |
| } |
| return mx::divide(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__rdiv__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("division", v); |
| } |
| return mx::divide(to_array(v, a.dtype()), a); |
| }, |
| "other"_a) |
| .def( |
| "__floordiv__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("floor division", v); |
| } |
| return mx::floor_divide(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__ifloordiv__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace floor division", v); |
| } |
| a.overwrite_descriptor(mx::floor_divide(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__rfloordiv__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("floor division", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| return mx::floor_divide(b, a); |
| }, |
| "other"_a) |
| .def( |
| "__mod__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("modulus", v); |
| } |
| return mx::remainder(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__imod__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace modulus", v); |
| } |
| a.overwrite_descriptor(mx::remainder(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__rmod__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("modulus", v); |
| } |
| return mx::remainder(to_array(v, a.dtype()), a); |
| }, |
| "other"_a) |
| .def( |
| "__eq__", |
| [](const mx::array& a, |
| const ScalarOrArray& v) -> std::variant<mx::array, bool> { |
| if (!is_comparable_with_array(v)) { |
| return false; |
| } |
| return mx::equal(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__lt__", |
| [](const mx::array& a, const ScalarOrArray v) -> mx::array { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("less than", v); |
| } |
| return mx::less(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__le__", |
| [](const mx::array& a, const ScalarOrArray v) -> mx::array { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("less than or equal", v); |
| } |
| return mx::less_equal(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__gt__", |
| [](const mx::array& a, const ScalarOrArray v) -> mx::array { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("greater than", v); |
| } |
| return mx::greater(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__ge__", |
| [](const mx::array& a, const ScalarOrArray v) -> mx::array { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("greater than or equal", v); |
| } |
| return mx::greater_equal(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__ne__", |
| [](const mx::array& a, |
| const ScalarOrArray v) -> std::variant<mx::array, bool> { |
| if (!is_comparable_with_array(v)) { |
| return true; |
| } |
| return mx::not_equal(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def("__neg__", [](const mx::array& a) { return -a; }) |
| .def("__bool__", [](mx::array& a) { return nb::bool_(to_scalar(a)); }) |
| .def( |
| "__repr__", |
| [](mx::array& a) { |
| nb::gil_scoped_release nogil; |
| std::ostringstream os; |
| os << a; |
| return os.str(); |
| }) |
| .def( |
| "__matmul__", |
| [](const mx::array& a, mx::array& other) { |
| return mx::matmul(a, other); |
| }, |
| "other"_a) |
| .def( |
| "__imatmul__", |
| [](mx::array& a, mx::array& other) -> mx::array& { |
| a.overwrite_descriptor(mx::matmul(a, other)); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__pow__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("power", v); |
| } |
| return mx::power(a, to_array(v, a.dtype())); |
| }, |
| "other"_a) |
| .def( |
| "__rpow__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("power", v); |
| } |
| return mx::power(to_array(v, a.dtype()), a); |
| }, |
| "other"_a) |
| .def( |
| "__ipow__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace power", v); |
| } |
| a.overwrite_descriptor(mx::power(a, to_array(v, a.dtype()))); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__invert__", |
| [](const mx::array& a) { |
| if (mx::issubdtype(a.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with bitwise inversion."); |
| } |
| if (a.dtype() == mx::bool_) { |
| return mx::logical_not(a); |
| } |
| return mx::bitwise_invert(a); |
| }) |
| .def( |
| "__and__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("bitwise and", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with bitwise and."); |
| } |
| return mx::bitwise_and(a, b); |
| }, |
| "other"_a) |
| .def( |
| "__iand__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace bitwise and", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with bitwise and."); |
| } |
| a.overwrite_descriptor(mx::bitwise_and(a, b)); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__or__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("bitwise or", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with bitwise or."); |
| } |
| return mx::bitwise_or(a, b); |
| }, |
| "other"_a) |
| .def( |
| "__ior__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace bitwise or", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with bitwise or."); |
| } |
| a.overwrite_descriptor(mx::bitwise_or(a, b)); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__lshift__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("left shift", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with left shift."); |
| } |
| return mx::left_shift(a, b); |
| }, |
| "other"_a) |
| .def( |
| "__ilshift__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace left shift", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with left shift."); |
| } |
| a.overwrite_descriptor(mx::left_shift(a, b)); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__rshift__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("right shift", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with right shift."); |
| } |
| return mx::right_shift(a, b); |
| }, |
| "other"_a) |
| .def( |
| "__irshift__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace right shift", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with right shift."); |
| } |
| a.overwrite_descriptor(mx::right_shift(a, b)); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def( |
| "__xor__", |
| [](const mx::array& a, const ScalarOrArray v) { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("bitwise xor", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed with bitwise xor."); |
| } |
| return mx::bitwise_xor(a, b); |
| }, |
| "other"_a) |
| .def( |
| "__ixor__", |
| [](mx::array& a, const ScalarOrArray v) -> mx::array& { |
| if (!is_comparable_with_array(v)) { |
| throw_invalid_operation("inplace bitwise xor", v); |
| } |
| auto b = to_array(v, a.dtype()); |
| if (mx::issubdtype(a.dtype(), mx::inexact) || |
| mx::issubdtype(b.dtype(), mx::inexact)) { |
| throw std::invalid_argument( |
| "Floating point types not allowed bitwise xor."); |
| } |
| a.overwrite_descriptor(mx::bitwise_xor(a, b)); |
| return a; |
| }, |
| "other"_a, |
| nb::rv_policy::none) |
| .def("__int__", [](mx::array& a) { return nb::int_(to_scalar(a)); }) |
| .def("__float__", [](mx::array& a) { return nb::float_(to_scalar(a)); }) |
| .def( |
| "__format__", |
| [](mx::array& a, nb::object format_spec) { |
| if (nb::len(nb::str(format_spec)) > 0 && a.ndim() > 0) { |
| throw nb::type_error( |
| "unsupported format string passed to mx.array.__format__"); |
| } else if (a.ndim() == 0) { |
| auto obj = to_scalar(a); |
| return nb::cast<std::string>( |
| nb::handle(PyObject_Format(obj.ptr(), format_spec.ptr()))); |
| } else { |
| nb::gil_scoped_release nogil; |
| std::ostringstream os; |
| os << a; |
| return os.str(); |
| } |
| }) |
| .def( |
| "flatten", |
| [](const mx::array& a, |
| int start_axis, |
| int end_axis, |
| const mx::StreamOrDevice& s) { |
| return mx::flatten(a, start_axis, end_axis, s); |
| }, |
| "start_axis"_a = 0, |
| "end_axis"_a = -1, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| R"pbdoc( |
| See :func:`flatten`. |
| )pbdoc") |
| .def( |
| "reshape", |
| [](const mx::array& a, nb::args shape_, mx::StreamOrDevice s) { |
| mx::Shape shape; |
| if (!nb::isinstance<int>(shape_[0])) { |
| shape = nb::cast<mx::Shape>(shape_[0]); |
| } else { |
| shape = nb::cast<mx::Shape>(shape_); |
| } |
| return mx::reshape(a, std::move(shape), s); |
| }, |
| "shape"_a, |
| "stream"_a = nb::none(), |
| R"pbdoc( |
| Equivalent to :func:`reshape` but the shape can be passed either as a |
| :obj:`tuple` or as separate arguments. |
| |
| See :func:`reshape` for full documentation. |
| )pbdoc") |
| .def( |
| "squeeze", |
| [](const mx::array& a, |
| const IntOrVec& v, |
| const mx::StreamOrDevice& s) { |
| if (std::holds_alternative<std::monostate>(v)) { |
| return mx::squeeze(a, s); |
| } else if (auto pv = std::get_if<int>(&v); pv) { |
| return mx::squeeze(a, *pv, s); |
| } else { |
| return mx::squeeze(a, std::get<std::vector<int>>(v), s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| R"pbdoc( |
| See :func:`squeeze`. |
| )pbdoc") |
| .def( |
| "abs", |
| &mx::abs, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`abs`.") |
| .def( |
| "__abs__", |
| [](const mx::array& a) { return mx::abs(a); }, |
| "See :func:`abs`.") |
| .def( |
| "square", |
| &mx::square, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`square`.") |
| .def( |
| "sqrt", |
| &mx::sqrt, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`sqrt`.") |
| .def( |
| "rsqrt", |
| &mx::rsqrt, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`rsqrt`.") |
| .def( |
| "reciprocal", |
| &mx::reciprocal, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`reciprocal`.") |
| .def( |
| "exp", |
| &mx::exp, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`exp`.") |
| .def( |
| "log", |
| &mx::log, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`log`.") |
| .def( |
| "log2", |
| &mx::log2, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`log2`.") |
| .def( |
| "log10", |
| &mx::log10, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`log10`.") |
| .def( |
| "sin", |
| &mx::sin, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`sin`.") |
| .def( |
| "cos", |
| &mx::cos, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`cos`.") |
| .def( |
| "log1p", |
| &mx::log1p, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`log1p`.") |
| .def( |
| "all", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::all(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`all`.") |
| .def( |
| "any", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::any(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`any`.") |
| .def( |
| "moveaxis", |
| &mx::moveaxis, |
| "source"_a, |
| "destination"_a, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`moveaxis`.") |
| .def( |
| "swapaxes", |
| &mx::swapaxes, |
| "axis1"_a, |
| "axis2"_a, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`swapaxes`.") |
| .def( |
| "transpose", |
| [](const mx::array& a, nb::args axes_, mx::StreamOrDevice s) { |
| if (axes_.size() == 0) { |
| return mx::transpose(a, s); |
| } |
| std::vector<int> axes; |
| if (!nb::isinstance<int>(axes_[0])) { |
| axes = nb::cast<std::vector<int>>(axes_[0]); |
| } else { |
| axes = nb::cast<std::vector<int>>(axes_); |
| } |
| return mx::transpose(a, axes, s); |
| }, |
| "axes"_a, |
| "stream"_a = nb::none(), |
| R"pbdoc( |
| Equivalent to :func:`transpose` but the axes can be passed either as |
| a tuple or as separate arguments. |
| |
| See :func:`transpose` for full documentation. |
| )pbdoc") |
| .def_prop_ro( |
| "T", |
| [](const mx::array& a) { return mx::transpose(a); }, |
| "Equivalent to calling ``self.transpose()`` with no arguments.") |
| .def( |
| "sum", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::sum(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`sum`.") |
| .def( |
| "prod", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::prod(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`prod`.") |
| .def( |
| "min", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::min(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`min`.") |
| .def( |
| "max", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::max(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`max`.") |
| .def( |
| "logcumsumexp", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool reverse, |
| bool inclusive, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::logcumsumexp(a, *axis, reverse, inclusive, s); |
| } else { |
| |
| |
| return mx::logcumsumexp( |
| mx::reshape(a, {-1}, s), 0, reverse, inclusive, s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| nb::kw_only(), |
| "reverse"_a = false, |
| "inclusive"_a = true, |
| "stream"_a = nb::none(), |
| "See :func:`logcumsumexp`.") |
| .def( |
| "logsumexp", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::logsumexp( |
| a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`logsumexp`.") |
| .def( |
| "mean", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| return mx::mean(a, get_reduce_axes(axis, a.ndim()), keepdims, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`mean`.") |
| .def( |
| "std", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| int ddof, |
| mx::StreamOrDevice s) { |
| return mx::std( |
| a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| "ddof"_a = 0, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`std`.") |
| .def( |
| "var", |
| [](const mx::array& a, |
| const IntOrVec& axis, |
| bool keepdims, |
| int ddof, |
| mx::StreamOrDevice s) { |
| return mx::var( |
| a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s); |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| "ddof"_a = 0, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`var`.") |
| .def( |
| "split", |
| [](const mx::array& a, |
| const std::variant<int, mx::Shape>& indices_or_sections, |
| int axis, |
| mx::StreamOrDevice s) { |
| if (auto pv = std::get_if<int>(&indices_or_sections); pv) { |
| return mx::split(a, *pv, axis, s); |
| } else { |
| return mx::split( |
| a, std::get<mx::Shape>(indices_or_sections), axis, s); |
| } |
| }, |
| "indices_or_sections"_a, |
| "axis"_a = 0, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`split`.") |
| .def( |
| "argmin", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::argmin(a, *axis, keepdims, s); |
| } else { |
| return mx::argmin(a, keepdims, s); |
| } |
| }, |
| "axis"_a = std::nullopt, |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`argmin`.") |
| .def( |
| "argmax", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool keepdims, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::argmax(a, *axis, keepdims, s); |
| } else { |
| return mx::argmax(a, keepdims, s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| "keepdims"_a = false, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`argmax`.") |
| .def( |
| "cumsum", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool reverse, |
| bool inclusive, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::cumsum(a, *axis, reverse, inclusive, s); |
| } else { |
| |
| |
| return mx::cumsum(reshape(a, {-1}, s), 0, reverse, inclusive, s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| nb::kw_only(), |
| "reverse"_a = false, |
| "inclusive"_a = true, |
| "stream"_a = nb::none(), |
| "See :func:`cumsum`.") |
| .def( |
| "cumprod", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool reverse, |
| bool inclusive, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::cumprod(a, *axis, reverse, inclusive, s); |
| } else { |
| |
| |
| return mx::cumprod( |
| mx::reshape(a, {-1}, s), 0, reverse, inclusive, s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| nb::kw_only(), |
| "reverse"_a = false, |
| "inclusive"_a = true, |
| "stream"_a = nb::none(), |
| "See :func:`cumprod`.") |
| .def( |
| "cummax", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool reverse, |
| bool inclusive, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::cummax(a, *axis, reverse, inclusive, s); |
| } else { |
| |
| |
| return mx::cummax( |
| mx::reshape(a, {-1}, s), 0, reverse, inclusive, s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| nb::kw_only(), |
| "reverse"_a = false, |
| "inclusive"_a = true, |
| "stream"_a = nb::none(), |
| "See :func:`cummax`.") |
| .def( |
| "cummin", |
| [](const mx::array& a, |
| std::optional<int> axis, |
| bool reverse, |
| bool inclusive, |
| mx::StreamOrDevice s) { |
| if (axis) { |
| return mx::cummin(a, *axis, reverse, inclusive, s); |
| } else { |
| |
| |
| return mx::cummin( |
| mx::reshape(a, {-1}, s), 0, reverse, inclusive, s); |
| } |
| }, |
| "axis"_a = nb::none(), |
| nb::kw_only(), |
| "reverse"_a = false, |
| "inclusive"_a = true, |
| "stream"_a = nb::none(), |
| "See :func:`cummin`.") |
| .def( |
| "round", |
| [](const mx::array& a, int decimals, mx::StreamOrDevice s) { |
| return mx::round(a, decimals, s); |
| }, |
| "decimals"_a = 0, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`round`.") |
| .def( |
| "diagonal", |
| [](const mx::array& a, |
| int offset, |
| int axis1, |
| int axis2, |
| mx::StreamOrDevice s) { |
| return mx::diagonal(a, offset, axis1, axis2, s); |
| }, |
| "offset"_a = 0, |
| "axis1"_a = 0, |
| "axis2"_a = 1, |
| "stream"_a = nb::none(), |
| "See :func:`diagonal`.") |
| .def( |
| "diag", |
| [](const mx::array& a, int k, mx::StreamOrDevice s) { |
| return mx::diag(a, k, s); |
| }, |
| "k"_a = 0, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| R"pbdoc( |
| Extract a diagonal or construct a diagonal matrix. |
| )pbdoc") |
| .def( |
| "conj", |
| [](const mx::array& a, mx::StreamOrDevice s) { |
| return mx::conjugate(to_array(a), s); |
| }, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`conj`.") |
| .def( |
| "view", |
| [](const ScalarOrArray& a, |
| const mx::Dtype& dtype, |
| mx::StreamOrDevice s) { return mx::view(to_array(a), dtype, s); }, |
| "dtype"_a, |
| nb::kw_only(), |
| "stream"_a = nb::none(), |
| "See :func:`view`."); |
| } |
|
|
