From e6035f576a639949cca78d56370664386b64feec Mon Sep 17 00:00:00 2001 From: Florian Fervers Date: Sun, 9 Jun 2024 22:10:48 +0200 Subject: [PATCH] Update docs --- docs/source/faq/backend.rst | 3 ++- docs/source/faq/universal.rst | 23 +++++++++---------- docs/source/index.rst | 1 + docs/source/more/backends.rst | 42 +++++++++++++++++++++++++++++++++++ 4 files changed, 56 insertions(+), 13 deletions(-) create mode 100644 docs/source/more/backends.rst diff --git a/docs/source/faq/backend.rst b/docs/source/faq/backend.rst index 05b44cd..882e020 100644 --- a/docs/source/faq/backend.rst +++ b/docs/source/faq/backend.rst @@ -2,7 +2,8 @@ How does einx support different tensor frameworks? ################################################## einx provides interfaces for tensor frameworks in the ``einx.backend.*`` namespace. einx functions accept a ``backend`` argument -that defines which backend to use for the computation. For ``backend=None`` (the default case), the backend is implicitly determined from the input tensors. +that defines which backend to use for the computation. For ``backend=None`` (the default case), the backend is implicitly determined +from the type of the input tensors. .. code:: python diff --git a/docs/source/faq/universal.rst b/docs/source/faq/universal.rst index 3470b94..0ef770f 100644 --- a/docs/source/faq/universal.rst +++ b/docs/source/faq/universal.rst @@ -34,15 +34,14 @@ einx notation einx provides an interface to tensor operations where vectorization is expressed entirely using einx notation, and each elementary operation is represented by exactly one API. The einx notation is: -* **Consistent**: The same type of notation is used for all elementary operations. -* **Unique**: Each elementary operation is represented by exactly one API. +* **Consistent**: The same type of notation is used for all elementary operations. Each elementary operation is represented by exactly one API. * **Complete**: Any operation that can be expressed with existing vectorization tools such as `jax.vmap `_ can also be expressed in einx notation. The following tables show examples of classical API calls that can be expressed using universal einx operations. -.. list-table:: Elementary operation: ``einx.get_at`` - :widths: 40 60 +.. list-table:: Example: ``einx.get_at`` + :widths: 42 58 :header-rows: 1 * - Classical API @@ -70,8 +69,8 @@ The following tables show examples of classical API calls that can be expressed | ``x[y[..., 0], y[..., 1]]`` - ``einx.get_at("a [...], a b [i] -> a b", x, y)`` -.. list-table:: Elementary operation: ``einx.dot`` (similar to einsum) - :widths: 40 60 +.. list-table:: Example: ``einx.dot`` (similar to einsum) + :widths: 42 58 :header-rows: 1 * - Classical API @@ -90,8 +89,8 @@ The following tables show examples of classical API calls that can be expressed * - ``np.inner(x, y)`` - ``einx.dot("x... [a], y... [a] -> x... y...", x, y)`` -.. list-table:: Elementary operation: ``einx.multiply`` - :widths: 40 60 +.. list-table:: Example: ``einx.multiply`` + :widths: 42 58 :header-rows: 1 * - Classical API @@ -105,8 +104,8 @@ The following tables show examples of classical API calls that can be expressed * - ``np.kron(x, y)`` - ``einx.multiply("a..., b... -> (a b)...", x, y),`` -.. list-table:: Elementary operation: ``einx.flip`` - :widths: 40 60 +.. list-table:: Example: ``einx.flip`` + :widths: 42 58 :header-rows: 1 * - Classical API @@ -114,9 +113,9 @@ The following tables show examples of classical API calls that can be expressed * - | ``np.flip(x, y, axis=0)`` | ``np.flipud(x, y)`` - - ``einx.flip("[a] b", x),`` + - ``einx.flip("[a] b", x)`` * - ``np.fliplr(x, y)`` - - ``einx.flip("a [b]", x),`` + - ``einx.flip("a [b]", x)`` .. * - ``einx.rearrange`` diff --git a/docs/source/index.rst b/docs/source/index.rst index ad648ed..ac1a86c 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -19,6 +19,7 @@ Welcome to einx's documentation! :maxdepth: 1 more/jit.rst + more/backends.rst more/gotchas.rst more/related.rst diff --git a/docs/source/more/backends.rst b/docs/source/more/backends.rst new file mode 100644 index 0000000..2cc8395 --- /dev/null +++ b/docs/source/more/backends.rst @@ -0,0 +1,42 @@ +Supported tensor frameworks +########################### + +**Numpy**: `https://numpy.org/ `_ + +Numpy does not support automatic vectorization (``vmap``). einx implements a custom ``vmap`` for Numpy instead using a Python for-loop +for testing and debugging purposes. This affects ``einx.vmap`` and ``einx.{index|get_at|set_at|...}``. + +---- + +**Torch**: `https://pytorch.org/ `_ + +einx disables ``torch.compile`` (using ``torch.compiler.disable``) when JIT-compiling a call into a Python function and reenables it when +executing the function. + +---- + +**Jax**: `https://jax.readthedocs.io/ `_ + +---- + +**Tensorflow**: `https://www.tensorflow.org/ `_ + +einx does not support tensors with dynamic shapes (i.e. ``None`` in the shape). + +---- + +**MLX**: `https://ml-explore.github.io/mlx `_ + +``einx.vmap`` and ``einx.{index|get_at|set_at|...}`` are currently not supported (``mx.vmap`` does not support all required primitives yet). + +---- + +**Tinygrad**: `https://tinygrad.org/ `_ + +``einx.vmap`` and ``einx.{index|get_at|set_at|...}`` are currently not supported. + +---- + +**Dask**: `https://docs.dask.org/en/stable/array.html `_ + +``einx.vmap`` and ``einx.{index|get_at|set_at|...}`` are currently not supported. \ No newline at end of file