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

Fix a few issues: docs for flatten, erf, dequantize validation (#560)

Browse Source 
* doc flatten

* erf doc

* check values for dequantize

* format
This commit is contained in:
Awni Hannun
2024-01-26 15:16:46 -08:00
committed by GitHub
parent bf17ab5002
commit 07f35c9d8a
2 changed files with 24 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/src/ops.cpp
View File

@@ -78,6 +78,11 @@ void init_ops(py::module_& m) {
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``.
@@ -87,6 +92,14 @@ void init_ops(py::module_& m) {
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(
"squeeze",
@@ -801,7 +814,7 @@ void init_ops(py::module_& m) {
Element-wise error function.
.. math::
\mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^t e^{-t^2} \, dx
\mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} \, dt
Args:
a (array): Input array.