lax.pdot() function status

[che-shr-cat]

Hi,

Could you please clarify the status of the lax.pdot() function?

It is absent in the documentation (https://jax.readthedocs.io/en/latest/jax.lax.html) and has no docstring at all (https://github.com/google/jax/blob/8d2cd5cce3377d59cc096a5e9c91f9856efc29a5/jax/_src/lax/parallel.py#L444).

However, it is used in some tests and examples, including the xmap() tutorial (https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html).

[jakevdp]

It was added in #5020 as an experimental primitive, and never documented. I suspect it was never intended to be a public function, though for historical reasons it is still exported publicly from jax.lax. @mattjj might be able to comment more on the status of this.

[mattjj]

We thought of it kind of like an analogue to lax.dot, in that it can be a public-facing function but a lower-level one, where people would likely instead prefer to use a jax.numpy-level wrapper. For lax.dot the natural jax.numpy wrapper is jnp.einsum. For lax.pdot, we expected jnp.einsum also to be the easier-to-use wrapper, since we extended the string spec syntax in jnp.einsum to handle named axes.

The reason we have wrappers at all is that the API that is easiest for humans to produce (i.e. to write) may not be the same as the API that is most natural for JAX's internal representation. The internal representation has other constraints on it, e.g. it tries to be amenable to transformations like autodiff and vmapping, tries to stick close to lower-level operations from XLA. More abstractly we could say that JAX's internal representations of operations need to be both easy to consume and easy to produce, in the sense that rules need to consume representations of these operations.

The lax.pdot function, like other functions introduced for xmap-style named axis work, hasn't gotten to the point of being well documented, since xmap itself is an experimental feature. Moreover work on this stuff has essentially been on pause in an incomplete state for a while. Instead we've been focused on automatic parallelization. I expect we'll revisit explicit named axis style parallelism this year, but we'll probably have a simpler take on it than xmap.

[che-shr-cat]

Thank you, @mattjj, for the clarification!