The current invocation of doctest from GH actions picked up only .rst files.
We enable .md files also, and we make a few changes to ensure that the
doctest passes on the existing files.
The changes fall into several categories:
* add a newline before the end of the code block, for doctest to
pick up the expected output properly
* update the expected values to match the current behavior
* disable some doctests that raise expected exceptions, whenever
I could not get doctest to match the exception details.
Sometimes +IGNORE_EXCEPTION_DETAIL was enough, and other times
I had to use +SKIP.
Fixed Typos in JEP doc files
Revert "Fixed Typos in JEP doc files"
This reverts commit c2a16950e0fc1b32971168501d183991e2394b5d.
revert two changes
reverted one change in advanced-autodiff
revert one change in parallelism
sync notebooks
JAX has not used mhlo for some time, in favor of stablehlo. Deprecate support for this dialect in JAX's API and remove testing.
PiperOrigin-RevId: 598550225
The new efficient-transpose path, enabled by setting check_rep=True in the shard_map call, had kept working. But the change inadvertently broke the check_rep=False path. And because most tests set check_rep=True, we didn't notice it in the tests!
The issue was that with check_rep=False, we need the shard_map transpose rule to insert psums corresponding to in_specs with fan-out, and correspondingly insert division for out_specs with fan-in-consensus. (With the new check_rep=True path that this change adds, those extra operations aren't necessary as the body itself transposes correctly.) But the PR accidentally removed those!
The fix was simple: just track whether we've applied the efficient-transpose-body-rewrite (i.e. whether we're in the new body-is-transposable path or old need-extra-operations path) by adding a boolean parameter `rewrite` to the shard_map primitive, and if the rewrite hasn't been applied then include the explicit psum/div operations in the transpose rule.
Reverts 8a04dfd830ff89f46e1fe3e866ee4fb2da9c90aa
PiperOrigin-RevId: 561805840
Currently
```
import jax
```
populates `jax._src` in the names exported from JAX. This change prepares for not exporting `jax._src` by default.
In particular, explicitly import modules from jax._src and refer to those imports rather than assuming jax._src contents will be around later. This is a common pattern in tests.
This change does not yet remove any exported names.
Issue https://github.com/google/jax/issues/11951
PiperOrigin-RevId: 469480816
