How is einx notation universal?

To address this question, let’s first look at how tensor operations are commonly expressed in existing tensor frameworks.

Classical notation#

Tensor operations can be dissected into two distinct components:

1. An elementary operation that is performed.

   • Example: np.sum computes a sum-reduction.

2. A division of the input tensor into sub-tensors. The elementary operation is applied to each sub-tensor independently. We refer to this as vectorization.

   • Example: Sub-tensors in np.sum span the dimensions specified by the axis parameter. The sum-reduction is vectorized over all other dimensions.

In common tensor frameworks like Numpy, PyTorch, Tensorflow or Jax, different elementary operations are implemented with different vectorization rules. For example, to express vectorization

• np.sum uses the axis parameter,

• np.add follows implicit broadcasting rules (e.g. in combination with np.newaxis), and

• np.matmul provides an implicit and custom set of rules.

Furthermore, an elementary operation is sometimes implemented in multiple APIs in order to offer vectorization rules for different use cases. For example, the retrieve-at-index operation can be implemented in PyTorch using tensor[coords], torch.gather, torch.index_select, torch.take, torch.take_along_dim, which conceptually apply the same low-level operation, but follow different vectorization rules (see below). Still, these interfaces sometimes do not cover all desirable use cases.

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. 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.

While elementary operations and vectorization are decoupled conceptually to provide a universal API, the implementation of the operations in the respective backend do not necessarily follow the same decoupling. For example, a matrix multiplication is represented as a vectorized dot-product in einx (using einx.dot), but still invokes an efficient matmul operation on the backend instead of a vectorized evaluation of the dot product.