// Copyright © 2023-2024 Apple Inc.

#include <numeric>

#include <ostream>

#include <variant>

#include <nanobind/nanobind.h>

#include <nanobind/stl/optional.h>

#include <nanobind/stl/pair.h>

#include <nanobind/stl/string.h>

#include <nanobind/stl/tuple.h>

#include <nanobind/stl/variant.h>

#include <nanobind/stl/vector.h>

#include "mlx/einsum.h"

#include "mlx/ops.h"

#include "mlx/utils.h"

#include "python/src/load.h"

#include "python/src/small_vector.h"

#include "python/src/utils.h"

namespace mx = mlx::core;

namespace nb = nanobind;

using namespace nb::literals;

using Scalar = std::variant<bool, int, double>;

mx::Dtype scalar_to_dtype(Scalar s) {

if (std::holds_alternative<int>(s)) {

return mx::int32;

} else if (std::holds_alternative<double>(s)) {

return mx::float32;

} else {

return mx::bool_;

}

}

double scalar_to_double(Scalar s) {

if (auto pv = std::get_if<int>(&s); pv) {

return static_cast<double>(*pv);

} else if (auto pv = std::get_if<double>(&s); pv) {

return *pv;

} else {

return static_cast<double>(std::get<bool>(s));

}

}

void init_ops(nb::module_& m) {

m.def(

"reshape",

&mx::reshape,

nb::arg(),

"shape"_a,

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def reshape(a: array, /, shape: Sequence[int], *, stream: "

"Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Reshape an array while preserving the size.

Args:

a (array): Input array.

shape (tuple(int)): New shape.

stream (Stream, optional): Stream or device. Defaults to ``None``

in which case the default stream of the default device is used.

Returns:

array: The reshaped array.

)pbdoc");

m.def(

"flatten",

[](const mx::array& a,

int start_axis,

int end_axis,

const mx::StreamOrDevice& s) {

return mx::flatten(a, start_axis, end_axis);

},

nb::arg(),

"start_axis"_a = 0,

"end_axis"_a = -1,

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def flatten(a: array, /, start_axis: int = 0, end_axis: int = "

"-1, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Flatten an array.

The axes flattened will be between ``start_axis`` and ``end_axis``,

inclusive. Negative axes are supported. After converting negative axis to

positive, axes outside the valid range will be clamped to a valid value,

``start_axis`` to ``0`` and ``end_axis`` to ``ndim - 1``.

Args:

a (array): Input array.

start_axis (int, optional): The first dimension to flatten. Defaults to ``0``.

end_axis (int, optional): The last dimension to flatten. Defaults to ``-1``.

stream (Stream, optional): Stream or device. Defaults to ``None``

in which case the default stream of the default device is used.

Returns:

array: The flattened array.

Example:

>>> a = mx.array([[1, 2], [3, 4]])

>>> mx.flatten(a)

array([1, 2, 3, 4], dtype=int32)

>>>

>>> mx.flatten(a, start_axis=0, end_axis=-1)

array([1, 2, 3, 4], dtype=int32)

)pbdoc");

m.def(

"unflatten",

&mx::unflatten,

nb::arg(),

"axis"_a,

"shape"_a,

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Unflatten an axis of an array to a shape.

Args:

a (array): Input array.

axis (int): The axis to unflatten.

shape (tuple(int)): The shape to unflatten to. At most one

entry can be ``-1`` in which case the corresponding size will be

inferred.

stream (Stream, optional): Stream or device. Defaults to ``None``

in which case the default stream of the default device is used.

Returns:

array: The unflattened array.

Example:

>>> a = mx.array([1, 2, 3, 4])

>>> mx.unflatten(a, 0, (2, -1))

array([[1, 2], [3, 4]], dtype=int32)

)pbdoc");

m.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);

}

},

nb::arg(),

"axis"_a = nb::none(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def squeeze(a: array, /, axis: Union[None, int, Sequence[int]] = "

"None, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Remove length one axes from an array.

Args:

a (array): Input array.

axis (int or tuple(int), optional): Axes to remove. Defaults

to ``None`` in which case all size one axes are removed.

Returns:

array: The output array with size one axes removed.

)pbdoc");

m.def(

"expand_dims",

[](const mx::array& a,

const std::variant<int, std::vector<int>>& v,

mx::StreamOrDevice s) {

if (auto pv = std::get_if<int>(&v); pv) {

return mx::expand_dims(a, *pv, s);

} else {

return mx::expand_dims(a, std::get<std::vector<int>>(v), s);

}

},

nb::arg(),

"axis"_a,

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def expand_dims(a: array, /, axis: Union[int, Sequence[int]], "

"*, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Add a size one dimension at the given axis.

Args:

a (array): Input array.

axes (int or tuple(int)): The index of the inserted dimensions.

Returns:

array: The array with inserted dimensions.

)pbdoc");

m.def(

"abs",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::abs(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def abs(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise absolute value.

Args:

a (array): Input array.

Returns:

array: The absolute value of ``a``.

)pbdoc");

m.def(

"sign",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::sign(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def sign(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise sign.

Args:

a (array): Input array.

Returns:

array: The sign of ``a``.

)pbdoc");

m.def(

"negative",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::negative(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def negative(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise negation.

Args:

a (array): Input array.

Returns:

array: The negative of ``a``.

)pbdoc");

m.def(

"add",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::add(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def add(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise addition.

Add two arrays with numpy-style broadcasting semantics. Either or both input arrays

can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The sum of ``a`` and ``b``.

)pbdoc");

m.def(

"subtract",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::subtract(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def subtract(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise subtraction.

Subtract one array from another with numpy-style broadcasting semantics. Either or both

input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The difference ``a - b``.

)pbdoc");

m.def(

"multiply",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::multiply(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def multiply(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise multiplication.

Multiply two arrays with numpy-style broadcasting semantics. Either or both

input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The multiplication ``a * b``.

)pbdoc");

m.def(

"divide",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::divide(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise division.

Divide two arrays with numpy-style broadcasting semantics. Either or both

input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The quotient ``a / b``.

)pbdoc");

m.def(

"divmod",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::divmod(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def divmod(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise quotient and remainder.

The fuction ``divmod(a, b)`` is equivalent to but faster than

``(a // b, a % b)``. The function uses numpy-style broadcasting

semantics. Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

tuple(array, array): The quotient ``a // b`` and remainder ``a % b``.

)pbdoc");

m.def(

"floor_divide",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::floor_divide(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def floor_divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise integer division.

If either array is a floating point type then it is equivalent to

calling :func:`floor` after :func:`divide`.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The quotient ``a // b``.

)pbdoc");

m.def(

"remainder",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::remainder(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def remainder(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise remainder of division.

Computes the remainder of dividing a with b with numpy-style

broadcasting semantics. Either or both input arrays can also be

scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The remainder of ``a // b``.

)pbdoc");

m.def(

"equal",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::equal(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise equality.

Equality comparison on two arrays with numpy-style broadcasting semantics.

Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The element-wise comparison ``a == b``.

)pbdoc");

m.def(

"not_equal",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::not_equal(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def not_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise not equal.

Not equal comparison on two arrays with numpy-style broadcasting semantics.

Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The element-wise comparison ``a != b``.

)pbdoc");

m.def(

"less",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::less(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def less(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise less than.

Strict less than on two arrays with numpy-style broadcasting semantics.

Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The element-wise comparison ``a < b``.

)pbdoc");

m.def(

"less_equal",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::less_equal(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def less_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise less than or equal.

Less than or equal on two arrays with numpy-style broadcasting semantics.

Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The element-wise comparison ``a <= b``.

)pbdoc");

m.def(

"greater",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::greater(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def greater(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise greater than.

Strict greater than on two arrays with numpy-style broadcasting semantics.

Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The element-wise comparison ``a > b``.

)pbdoc");

m.def(

"greater_equal",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::greater_equal(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def greater_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise greater or equal.

Greater than or equal on two arrays with numpy-style broadcasting semantics.

Either or both input arrays can also be scalars.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The element-wise comparison ``a >= b``.

)pbdoc");

m.def(

"array_equal",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

bool equal_nan,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::array_equal(a, b, equal_nan, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"equal_nan"_a = false,

"stream"_a = nb::none(),

nb::sig(

"def array_equal(a: Union[scalar, array], b: Union[scalar, array], equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Array equality check.

Compare two arrays for equality. Returns ``True`` if and only if the arrays

have the same shape and their values are equal. The arrays need not have

the same type to be considered equal.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

equal_nan (bool): If ``True``, NaNs are considered equal.

Defaults to ``False``.

Returns:

array: A scalar boolean array.

)pbdoc");

m.def(

"matmul",

&mx::matmul,

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def matmul(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Matrix multiplication.

Perform the (possibly batched) matrix multiplication of two arrays. This function supports

broadcasting for arrays with more than two dimensions.

- If the first array is 1-D then a 1 is prepended to its shape to make it

a matrix. Similarly if the second array is 1-D then a 1 is appended to its

shape to make it a matrix. In either case the singleton dimension is removed

from the result.

- A batched matrix multiplication is performed if the arrays have more than

2 dimensions. The matrix dimensions for the matrix product are the last

two dimensions of each input.

- All but the last two dimensions of each input are broadcast with one another using

standard numpy-style broadcasting semantics.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The matrix product of ``a`` and ``b``.

)pbdoc");

m.def(

"square",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::square(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def square(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise square.

Args:

a (array): Input array.

Returns:

array: The square of ``a``.

)pbdoc");

m.def(

"sqrt",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::sqrt(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def sqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise square root.

Args:

a (array): Input array.

Returns:

array: The square root of ``a``.

)pbdoc");

m.def(

"rsqrt",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::rsqrt(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def rsqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise reciprocal and square root.

Args:

a (array): Input array.

Returns:

array: One over the square root of ``a``.

)pbdoc");

m.def(

"reciprocal",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::reciprocal(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def reciprocal(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise reciprocal.

Args:

a (array): Input array.

Returns:

array: The reciprocal of ``a``.

)pbdoc");

m.def(

"logical_not",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::logical_not(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def logical_not(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise logical not.

Args:

a (array): Input array or scalar.

Returns:

array: The boolean array containing the logical not of ``a``.

)pbdoc");

m.def(

"logical_and",

[](const ScalarOrArray& a, const ScalarOrArray& b, mx::StreamOrDevice s) {

return mx::logical_and(to_array(a), to_array(b), s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def logical_and(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise logical and.

Args:

a (array): First input array or scalar.

b (array): Second input array or scalar.

Returns:

array: The boolean array containing the logical and of ``a`` and ``b``.

)pbdoc");

m.def(

"logical_or",

[](const ScalarOrArray& a, const ScalarOrArray& b, mx::StreamOrDevice s) {

return mx::logical_or(to_array(a), to_array(b), s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def logical_or(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise logical or.

Args:

a (array): First input array or scalar.

b (array): Second input array or scalar.

Returns:

array: The boolean array containing the logical or of ``a`` and ``b``.

)pbdoc");

m.def(

"logaddexp",

[](const ScalarOrArray& a_,

const ScalarOrArray& b_,

mx::StreamOrDevice s) {

auto [a, b] = to_arrays(a_, b_);

return mx::logaddexp(a, b, s);

},

nb::arg(),

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def logaddexp(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise log-add-exp.

This is a numerically stable log-add-exp of two arrays with numpy-style

broadcasting semantics. Either or both input arrays can also be scalars.

The computation is is a numerically stable version of ``log(exp(a) + exp(b))``.

Args:

a (array): Input array or scalar.

b (array): Input array or scalar.

Returns:

array: The log-add-exp of ``a`` and ``b``.

)pbdoc");

m.def(

"exp",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::exp(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def exp(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise exponential.

Args:

a (array): Input array.

Returns:

array: The exponential of ``a``.

)pbdoc");

m.def(

"expm1",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::expm1(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def expm1(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise exponential minus 1.

Computes ``exp(x) - 1`` with greater precision for small ``x``.

Args:

a (array): Input array.

Returns:

array: The expm1 of ``a``.

)pbdoc");

m.def(

"erf",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::erf(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def erf(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise error function.

.. math::

\mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} \, dt

Args:

a (array): Input array.

Returns:

array: The error function of ``a``.

)pbdoc");

m.def(

"erfinv",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::erfinv(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def erfinv(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise inverse of :func:`erf`.

Args:

a (array): Input array.

Returns:

array: The inverse error function of ``a``.

)pbdoc");

m.def(

"sin",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::sin(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def sin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise sine.

Args:

a (array): Input array.

Returns:

array: The sine of ``a``.

)pbdoc");

m.def(

"cos",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::cos(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def cos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise cosine.

Args:

a (array): Input array.

Returns:

array: The cosine of ``a``.

)pbdoc");

m.def(

"tan",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::tan(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def tan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise tangent.

Args:

a (array): Input array.

Returns:

array: The tangent of ``a``.

)pbdoc");

m.def(

"arcsin",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {

return mx::arcsin(to_array(a), s);

},

nb::arg(),

nb::kw_only(),

"stream"_a = nb::none(),

nb::sig(

"def arcsin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),

R"pbdoc(

Element-wise inverse sine.

Args:

a (array): Input array.

Returns:

array: The inverse sine of ``a``.

)pbdoc");

m.def(

"arccos",

[](const ScalarOrArray& a, mx::StreamOrDevice s) {