diff --git a/docs/src/usage/function_transforms.rst b/docs/src/usage/function_transforms.rst index 9a15bbf1..9769fcea 100644 --- a/docs/src/usage/function_transforms.rst +++ b/docs/src/usage/function_transforms.rst @@ -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)) 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: @@ -169,7 +169,7 @@ Instead you can use :func:`vmap` to automatically vectorize the addition: # Vectorize over the second dimension of x and the # 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 corresponding input to vectorize over. Similarly, use ``out_axes`` to specify diff --git a/docs/src/usage/indexing.rst b/docs/src/usage/indexing.rst index 62994a0f..c74e357f 100644 --- a/docs/src/usage/indexing.rst +++ b/docs/src/usage/indexing.rst @@ -77,7 +77,7 @@ from the GPU. Performing bounds checking for array indices before launching the kernel would be extremely inefficient. 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 operations which MLX does not yet support include :func:`numpy.nonzero` and the single input version of :func:`numpy.where`.