From bed906390d8cd3e6c0e4edcefcb1377fcd6caaf3 Mon Sep 17 00:00:00 2001 From: Florian Fervers Date: Thu, 20 Jun 2024 09:40:47 +0200 Subject: [PATCH] Update docs --- docs/source/gettingstarted/tutorial_ops.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/gettingstarted/tutorial_ops.rst b/docs/source/gettingstarted/tutorial_ops.rst index 4350023..d1f6299 100644 --- a/docs/source/gettingstarted/tutorial_ops.rst +++ b/docs/source/gettingstarted/tutorial_ops.rst @@ -279,7 +279,7 @@ In :func:`einx.vmap`, the input and output tensors of ``op`` match the marked ax einx.vmap("a [b c] -> a [b]", x, op=op) -:func:`einx.vmap` is implemented using efficient automatic vectorization in the respective backend (e.g. +:func:`einx.vmap` is implemented using automatic vectorization in the respective backend (e.g. `jax.vmap `_, `torch.vmap `_). einx also implements a simple ``vmap`` function for the Numpy backend for testing/ debugging purposes using a Python loop. @@ -293,7 +293,7 @@ In :func:`einx.vmap_with_axis`, ``op`` is instead given an ``axis`` argument and # Input: x has shape "a b c", axis is (1, 2) x = np.sum(x, axis=axis[1]) x = np.flip(x, axis=axis[0]) - # Output: x has shape "b" + # Output: x has shape "a b" return x einx.vmap_with_axis("(a [b c]) -> (a [b])", x, op=op, a=2, b=3, c=4)