| | |
| | #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`."); |
| | } |
| |
|
