zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix vmap example in docs (#1556)

Browse Source
This commit is contained in:
Chris Offner 2024-11-03 01:44:14 +01:00 committed by GitHub
parent 42533931fa
commit 46d8b16ab4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docs/src/usage/function_transforms.rst
View File

@ -161,7 +161,7 @@ A naive way to add the elements from two sets of vectors is with a loop:
ys = mx.random.uniform(shape=(100, 4096)) ys = mx.random.uniform(shape=(100, 4096))
def naive_add(xs, ys): def naive_add(xs, ys):
return [xs[i] + ys[:, i] for i in range(xs.shape[1])] return [xs[i] + ys[:, i] for i in range(xs.shape[0])]
Instead you can use :func:`vmap` to automatically vectorize the addition: Instead you can use :func:`vmap` to automatically vectorize the addition:
@ -169,7 +169,7 @@ Instead you can use :func:`vmap` to automatically vectorize the addition:
# Vectorize over the second dimension of x and the # Vectorize over the second dimension of x and the
# first dimension of y # first dimension of y
vmap_add = mx.vmap(lambda x, y: x + y, in_axes=(1, 0)) vmap_add = mx.vmap(lambda x, y: x + y, in_axes=(0, 1))
The ``in_axes`` parameter can be used to specify which dimensions of the The ``in_axes`` parameter can be used to specify which dimensions of the
corresponding input to vectorize over. Similarly, use ``out_axes`` to specify corresponding input to vectorize over. Similarly, use ``out_axes`` to specify

2
docs/src/usage/indexing.rst
View File

@ -77,7 +77,7 @@ from the GPU. Performing bounds checking for array indices before launching the
kernel would be extremely inefficient. kernel would be extremely inefficient.
Indexing with boolean masks is something that MLX may support in the future. In Indexing with boolean masks is something that MLX may support in the future. In
general, MLX has limited support for operations for which outputs general, MLX has limited support for operations for which output
*shapes* are dependent on input *data*. Other examples of these types of *shapes* are dependent on input *data*. Other examples of these types of
operations which MLX does not yet support include :func:`numpy.nonzero` and the operations which MLX does not yet support include :func:`numpy.nonzero` and the
single input version of :func:`numpy.where`. single input version of :func:`numpy.where`.