This means that jax and its dependencies (e.g. jaxlib) must be
manually installed before running the tests. This is useful for
testing an existing jax install, e.g. a later version of jaxlib, GPU
jaxlib, etc.
JAX_NUM_GENERATED_CASES=25 is now using kokoro with py3.9 `CPU` which is 4x faster than the github action.
Since py3.8 build with JAX_NUM_GENERATED_CASES=25 on github action is the slowest build, this should give an improvement of 10 minutes overall since the second slowest build is 10 mins faster than the above one.
But moving py3.8 build (Github actions) with JAX_NUM_GENERATED_CASES=25 to Kokoro (JAX_NUM_GENERATED_CASES=25 and py3.9 and `CPU`) gives a `4x` speedup (testing time).
Overall kokoro time is around 11 mins. So its a `2.5x` speedup overall.
Github actions:
```
=============== 23033 passed, 2033 skipped in 1727.43s (0:28:47) ===============
```
Kokoro:
```
=============== 22877 passed, 2065 skipped in 446.36s (0:07:26) ================
Total kokoro time = 11 mins. So 2.5x speedup overall.
```
PiperOrigin-RevId: 390656668
👋 hello there! I'm a fellow Googler who works on projects that leverage GitHub Actions for CI/CD. Recently I noticed a large increase in our queue time, and I've tracked it down to the [limit of 180 concurrent jobs](https://docs.github.com/en/actions/reference/usage-limits-billing-and-administration) for an organization. To help be better citizens, I'm proposing changes across a few repositories that will reduce GitHub Actions hours and consumption. I hope these changes are reasonable and I'm happy to talk through them in more detail.
- **(you were already doing this, thank you!**) Only run GitHub Actions for pushes and PRs against the main branch of the repository. If your team uses a forking model, this change will not affect you. If your team pushes branches to the repository directly, this changes actions to only run against the primary branches or if you open a Pull Request against a primary branch.
- For long-running jobs (especially tests), I added the "Cancel previous" workflow. This is very helpful to prevent a large queue backlog when you are doing rapid development and pushing multiple commits. Without this, GitHub Actions' default behavior is to run all actions on all commits.
There are other changes you could make, depending on your project (but I'm not an expert):
- If you have tests that should only run when a subset of code changes, consider gating your workflow to particular file paths. For example, we have some jobs that do Terraform linting, but [they only run when Terraform files are changed](c4f59fee71/.github/workflows/terraform.yml (L3-L11)).
Hopefully these changes are not too controversial and also hopefully you can see how this would reduce actions consumption to be good citizens to fellow Googlers. If you have any questions, feel free to respond here or ping me on chat. Thank you!
Revert also previous changes that pinned numpy to 1.19.
One of the changes in numpy 1.20 is to add more type annotations.
However, this sometimes make mypy give errors. A common example is
numpy.take, which with the new type annotation does not appear to
mypy as indexable.
Another change is that np.int and np.bool are deprecated. One
should use np.bool_ or np.int_, or the built-ins bool and int.
As numpy 1.20 was released recently, and triggers some errors
in the GitHub CI, we pin numpy to 1.19. It seems that we still
get failures when trying to import numpy-dispatch. We
disable it until we figure out the problem.
* Add experimental __array_module__ method
xref https://github.com/google/jax/issues/1565
`__array_module__` (see [NEP 37](https://numpy.org/neps/nep-0037-array-module.html))
is an experimental alternative to `__array_function__` and `__array_ufunc__`
for "duck array" compatibility with NumPy that promises to be much less
invasive.
Example usage:
```python
import numpy as np
def duckarray_stack(arrays):
"""This "stack" function should work with any array library, including JAX."""
npx = np.get_array_module(*arrays)
arrays = [npx.asarray(arr) for arr in arrays]
shapes = {arr.shape for arr in arrays}
if len(shapes) != 1:
raise ValueError('all input arrays must have the same shape')
expanded_arrays = [arr[npx.newaxis, ...] for arr in arrays]
return npx.concatenate(expanded_arrays, axis=0)
```
Support for this protocol has *not* yet been implemented in NumPy, but it can
be tested with https://github.com/seberg/numpy-dispatch.
My reasoning for merging it into JAX (on an experimental basis with no
guarantees, of course) is that:
1. It's not invasive -- the implementation is small and self-contained.
2. No backwards compatibility issues. Unlike `__array_function__` and
`__array_ufunc__`, `__array_module__` will always require an explicit
opt-in by libraries that use it by calling `get_array_module()`.
2. Other NumPy developers
[want evidence](https://github.com/numpy/numpy/pull/16935#issuecomment-673951287)
that this is actually feasible.
3. Scikit-Learn developers like @thomasjpfan are interested in exploring
supporting scikit-learn on top of NumPy-like libraries like JAX, and
experimental support for this protocol will make that easier.
Note: this PR does add `numpy-dispatch` as a optional testing requirement in
order to verify that this works. If desired, we could remove this from CI, but
installing numpy-dispatch (and its build requirement Cython) appears to only
add a few seconds of build time.
* don't explicitly list cython
* remove UnshpaedArray from _JAX_ARRAY_TYPES
* Remove incorrect note about metaclasses
* remove unnecessary numpy_dispatch.ensure_dispatching()
We do this to see if this reduces the incidence of errors fetching
the tf-nightly package. These tests are being run when we import
the code in Google.
This change, when enabled, stages out all primitive calls in the dynamic
scope of a jitted, pmapped, or control flow function, rather than only
staging out based on data dependence. One improvement is that jitted
functions can consume less memory, by avoiding instantiating large
constants at trace time, and cause less memory fragmentation as well. It
also simplifies several internals.
See https://github.com/google/jax/pull/3370 fo more information.
* Refactored host_callback to use the C++ runtime.
* The new runtime makes it unnecessary to start the outfeed_receiver
in the user's code
* We don't need msgpack anymore
* There is an interaction between host_callback and using lax.outfeed.
I am trying to solve this by (a) making host_callback_test stop the
outfeed receiver on finish and infeed_test on start, and (b)
telling pytest-xdist to run all the tests from one file into
a single worker.
