Repository: google/jax

Issue:
jax.lax.dot_general has minimal documentation
Please:

- [x] Check for duplicate requests.
- [x] Describe your goal, and if possible provide a code snippet with a motivating example.

----------

The docs for [`jax.lax.dot_general`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dot_general.html) are unclear. In particular, the `dimension_numbers` argument isn't clear to me.


Discussion:
Thanks for the request!

Just brainstorming ways to improve: do the [XLA HLO docs](https://www.tensorflow.org/xla/operation_semantics#dotgeneral) clarify the `dimension_numbers` argument? Or are those unclear too?

There is a [reference implementation in terms of standard `numpy`](https://github.com/google/jax/blob/d27735820049bdb9b0828ee95abf86d1ec50f1f9/jax/_src/lax_reference.py#L208) which we could point to, but I don't think it's concise/clear enough to be helpful.

(By the way, I went to elementary school with someone with your name... any chance you're from The Peninsula originally?)
I wonder if it might make sense to add a note similar to what we have in [`lax.gather`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.gather.html), something along the these lines:

> The semantics of `dot_general` are complicated, but most users should not have to use it directly. Instead, you can use more user-friendly functions like `jax.numpy.dot`, `jax.numpy.tensordot`, `jax.numpy.einsum`, and others which will construct appropriate calls to `dot_general` under the hood. If you really want to understand `dot_general` itself, we recommend reading XLA's [DotGeneral](https://www.tensorflow.org/xla/operation_semantics#dotgeneral) documentation.
@jakevdp A note like that would be helpful.

@mattjj The XLA docs are more helpful, but still not totally clear. This line is the most helpful

> It follows that the resulting dimension number starts with the batch dimension, then the 'lhs' non-contracting/non-batch dimension, and finally the 'rhs' non-contracting/non-batch dimension.

However, "It follows that" makes it unclear whether it's implied by something about the example, or if it's general to the function. In the jax docs, it might be sufficient to simply specify output dimensions in terms of the inputs. If it's a general property of the return, a change like this could improve things greatly:

> Returns: An array whose first dimensions are the batch dimensions, followed by the the 'lhs' non-contracting/non-batch dimensions, and finally the 'rhs' non-contracting/non-batch dimensions.
Proposed doc update in https://github.com/google/jax/pull/14674; let me know what you think!

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -747,18 +747,25 @@ def dot(lhs: Array, rhs: Array, precision: PrecisionLike = None,
 def dot_general(lhs: ArrayLike, rhs: ArrayLike, dimension_numbers: DotDimensionNumbers,
                 precision: PrecisionLike = None,
                 preferred_element_type: Optional[DTypeLike] = None) -> Array:
-  """More general contraction operator.
+  """General dot product/contraction operator.
 
   Wraps XLA's `DotGeneral
   <https://www.tensorflow.org/xla/operation_semantics#dotgeneral>`_
   operator.
 
+  The semantics of ``dot_general`` are complicated, but most users should not have to
+  use it directly. Instead, you can use higher-level functions like :func:`jax.numpy.dot`,
+  :func:`jax.numpy.matmul`, :func:`jax.numpy.tensordot`, :func:`jax.numpy.einsum`,
+  and others which will construct appropriate calls to ``dot_general`` under the hood.
+  If you really want to understand ``dot_general`` itself, we recommend reading XLA's
+  `DotGeneral  <https://www.tensorflow.org/xla/operation_semantics#dotgeneral>`_
+  operator documentation.
+
   Args:
     lhs: an array
     rhs: an array
-    dimension_numbers: a tuple of tuples of the form
-      `((lhs_contracting_dims, rhs_contracting_dims),
-      (lhs_batch_dims, rhs_batch_dims))`
+    dimension_numbers: a tuple of tuples of sequences of ints of the form
+      ``((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims))``
     precision: Optional. Either ``None``, which means the default precision for
       the backend, a :class:`~jax.lax.Precision` enum value (``Precision.DEFAULT``,
       ``Precision.HIGH`` or ``Precision.HIGHEST``) or a tuple of two
@@ -768,7 +775,9 @@ def dot_general(lhs: ArrayLike, rhs: ArrayLike, dimension_numbers: DotDimensionN
       accumulate results to and return a result with that datatype.
 
   Returns:
-    An array containing the result.
+    An array whose first dimensions are the (shared) batch dimensions, followed by
+    the ``lhs`` non-contracting/non-batch dimensions, and finally the ``rhs``
+    non-contracting/non-batch dimensions.
   """
   (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers
   cdims = (api_util._ensure_index_tuple(lhs_contract),


Test patch: