zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 10:26:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

std and expm1 (#973)

Browse Source 
* std and expm1

* actually add expm1

* fix linux

* fix vjp

* relax tol for linux test

* Add it to the compilable primitives

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Awni Hannun
2024-04-08 14:26:01 -07:00
committed by GitHub
parent 76e63212ff
commit 42afe27e12
19 changed files with 332 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
python/src/ops.cpp
View File

@@ -772,6 +772,25 @@ void init_ops(nb::module_& m) {
Returns:
array: The exponential of ``a``.
)pbdoc");
m.def(
"expm1",
&mlx::core::expm1,
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",
&mlx::core::erf,
@@ -2150,6 +2169,40 @@ void init_ops(nb::module_& m) {
Returns:
array: The output array of variances.
)pbdoc");
m.def(
"std",
[](const array& a,
const IntOrVec& axis,
bool keepdims,
int ddof,
StreamOrDevice s) {
return mlx::core::std(
a, get_reduce_axes(axis, a.ndim()), keepdims, ddof, s);
},
nb::arg(),
"axis"_a = nb::none(),
"keepdims"_a = false,
"ddof"_a = 0,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def std(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Compute the standard deviation(s) over the given axes.
Args:
a (array): Input array.
axis (int or list(int), optional): Optional axis or
axes to reduce over. If unspecified this defaults
to reducing over the entire array.
keepdims (bool, optional): Keep reduced axes as
singleton dimensions, defaults to `False`.
ddof (int, optional): The divisor to compute the variance
is ``N - ddof``, defaults to 0.
Returns:
array: The output array of standard deviations.
)pbdoc");
m.def(
"split",
[](const array& a,