Repository: google/jax

Issue:
jnp.einsum docstring missing equation arg
### Description

The documented arg list for einsum is missing the first argument.

https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.einsum.html
`jax.numpy.einsum(*operands, out=None, optimize='optimal', precision=None, _use_xeinsum=False)`
compare to
https://numpy.org/doc/stable/reference/generated/numpy.einsum.html
`numpy.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)`
(noting the 'subscripts' for the first arg)

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Hi - thanks for the report! It looks like the issue is that the actual signature of `np.einsum` doesn't match its documented signature. For example:
```python
In [1]: import numpy as np

In [2]: np.einsum??
Signature: np.einsum(*operands, out=None, optimize=False, **kwargs)
Source:   
@array_function_dispatch(_einsum_dispatcher, module='numpy')
def einsum(*operands, out=None, optimize=False, **kwargs):
    """
    einsum(subscripts, *operands, out=None, dtype=None, order='K',
           casting='safe', optimize=False)

    Evaluates the Einstein summation convention on the operands.
...
```
Note that the coded signature matches the JAX implementation, but the signature written at the top of the docstring is different. Given that, I think JAX's implementation is correct, because it matches numpy's implementation. What do you think?
I suspect the reason for the mismatch on numpy's side is that the signature pre-dates python 3.0, so keyword-only arguments weren't originally possible to use here.
The numpy documentation is more helpful if you don't have memorized that
the equation comes first. An example usage would be even better.

On Wed, Mar 15, 2023, 6:57 PM Jake Vanderplas ***@***.***>
wrote:

> I suspect the reason for the mismatch on numpy's side is that the
> signature pre-dates python 3.0, so keyword-only arguments weren't
> originally possible to use here.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/google/jax/issues/15017#issuecomment-1470950228>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AFJFSI2NIYVEUTAABZLYDBTW4JCLZANCNFSM6AAAAAAV4L6S5Q>
> .
> You are receiving this because you authored the thread.Message ID:
> ***@***.***>
>

Thanks - for `jax.numpy` functions, the JAX documentation is just a direct copy of the numpy documentation in most cases.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3112,6 +3112,7 @@ def tensordot(a, b, axes=2, *, precision=None):
 
 @util._wraps(np.einsum, lax_description=_EINSUM_DOC, skip_params=['out'])
 def einsum(
+    subscripts,
     *operands,
     out=None,
     optimize="optimal",
@@ -3119,6 +3120,7 @@ def einsum(
     _use_xeinsum=False,
     _dot_general=lax.dot_general,
 ):
+  operands = (subscripts, *operands)
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.einsum is not supported.")
 


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -5251,7 +5251,7 @@ def testWrappedSignaturesMatch(self):
     }
 
     extra_params = {
-      'einsum': ['precision'],
+      'einsum': ['subscripts', 'precision'],
       'einsum_path': ['subscripts'],
       'take_along_axis': ['mode'],
     }