7287 Commits

Author SHA1 Message Date
Skye Wanderman-Milne
5cb5056ea7 Suppress gRPC log spam on Cloud TPU. 2021-03-11 22:52:54 +00:00
jax authors
077793cd64 Merge pull request #6019 from skye:examples_test
PiperOrigin-RevId: 362363477
2021-03-11 13:23:53 -08:00
jax authors
65ee6041cb Merge pull request #6020 from jamestwebber:patch-1
PiperOrigin-RevId: 362360958
2021-03-11 13:11:21 -08:00
James Webber
cf4f44548c
sync markdown 2021-03-11 14:24:41 -05:00
jax authors
6b923a9850 Merge pull request #6026 from Chillee:patch-1
PiperOrigin-RevId: 362328216
2021-03-11 10:48:43 -08:00
Horace He
abdee43688
Update 04-advanced-autodiff.md 2021-03-11 05:56:52 -05:00
Horace He
e38d272357
Update 04-advanced-autodiff.ipynb 2021-03-11 05:55:48 -05:00
Jasmijn Bastings
6470094f04
Fix import of datasets in differentially private SGD example 2021-03-11 11:32:43 +01:00
jax authors
7ef86ef8ad Merge pull request #6013 from zhangqiaorjc:new_dlpack_fix
PiperOrigin-RevId: 362224066
2021-03-10 23:19:44 -08:00
Skye Wanderman-Milne
cd619978bb Adjust precision in examples/kernel_lsq.py and corresponding test.
This is important on TPU.
2021-03-11 03:03:43 +00:00
James Webber
3a99e58033
(hopefully) fix broken link 2021-03-10 20:09:54 -05:00
James Bradbury
72a3036b1a Hotfix for another assertion that's too strict about named shapes
PiperOrigin-RevId: 362164157
2021-03-10 16:09:37 -08:00
Seth Vargo
9b7aa37a12
Update ci-build.yaml 2021-03-10 17:55:21 -05:00
jax authors
cf9b77f1de Merge pull request #5998 from zhangqiaorjc:dev_put_count
PiperOrigin-RevId: 362143966
2021-03-10 14:36:55 -08:00
Qiao Zhang
e523ae0206 Add a test that fails due to unsupported strides.
- test fails during dlpack tensor to buffer conversion
- XLA layout does not support arbitrary dlpack strides
- users should explicit materialze such tensors by making a copy
2021-03-10 14:19:21 -08:00
Matthew Johnson
ecbc13475c fix typo 2021-03-10 14:16:48 -08:00
Seth Vargo
1048ebb9fa
Update ci-build.yaml 2021-03-10 17:14:48 -05:00
Seth Vargo
e550dda8ba
Improve Actions hygiene
👋 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!
2021-03-10 16:26:01 -05:00
jax authors
61041cb1e3 Merge pull request #6010 from hawkinsp:issue4690
PiperOrigin-RevId: 362120286
2021-03-10 12:49:55 -08:00
Peter Hawkins
62a726d329 Add workaround for SelectAndScatter padding bug on CPU and GPU. 2021-03-10 15:25:32 -05:00
Matthew Johnson
d11bba9caf tweak broken tests 2021-03-10 12:20:36 -08:00
jax authors
c9c89c4820 Merge pull request #5997 from jakevdp:fix-piecewise
PiperOrigin-RevId: 362085337
2021-03-10 10:32:56 -08:00
jax authors
b675905db8 Merge pull request #6012 from hawkinsp:lazy3
PiperOrigin-RevId: 362079742
2021-03-10 10:14:08 -08:00
jax authors
d0de749135 Merge pull request #5996 from skye:cudnn_versions
PiperOrigin-RevId: 362077738
2021-03-10 10:06:21 -08:00
jax authors
5098caf1fc Merge pull request #5984 from skye:cloud_tpu_env_vars
PiperOrigin-RevId: 362072957
2021-03-10 09:49:40 -08:00
Skye Wanderman-Milne
c32d1e5aae Automatically initialize Cloud TPU topology env vars if running on a Cloud TPU VM.
This removes the need to manually set these env vars when running on a Cloud TPU pod slice.
2021-03-10 09:15:31 -08:00
jax authors
0acafcf8c0 Merge pull request #6011 from hawkinsp:token
PiperOrigin-RevId: 362060386
2021-03-10 08:59:12 -08:00
Peter Hawkins
5d6ff8b28d Readd omnistaging friendly versions of large constant tests. 2021-03-10 11:34:42 -05:00
Peter Hawkins
00349390fe Fix crash returning a Token from a jit computation on GPU.
Calling .numpy_dtype() doesn't work on tokens. But we don't need a numpy dtype here, an XLA dtype works just as well.
2021-03-10 10:18:38 -05:00
jax authors
3e45a8376c Merge pull request #6002 from hawkinsp:lazy3
PiperOrigin-RevId: 362034533
2021-03-10 06:27:20 -08:00
Sharad Vikram
ddaef193fe Add scan and while rule for jax.experimental.callback transformation 2021-03-09 19:46:16 -08:00
jax authors
23b82b99d2 Merge pull request #6001 from skye:jaxlib_version
PiperOrigin-RevId: 361948443
2021-03-09 18:56:54 -08:00
Peter Hawkins
140c0acbbe Remove the JAX lazy sublanguage.
Back in the mists of time, before omnistaging landed in JAX, we used lazy
expressions to avoid materializing large constants inside `jit` computations.
Omnistaging, which means that computations that are in the dynamic scope of a
`jit` are staged into the `jit` computation, has subsumed most of the reasons
for laziness to exist, and this PR removes the laziness support for simplicity.

At the time of this PR, laziness is used only for broadcasts and transposes in
eager mode (i.e., outside a `jit`). This allows us to:
a) fuse together multiple broadcasts and transposes, and
b) if a lazy expression is lexically captured by a `jit` computation, we can
   avoid materializing it in its expanded form.

It is not clear that laziness has sufficient power to weight ratio to continue
to exist, and it is making other work on improving JAX dispatch times more
difficult. As a result, this PR removes laziness to unblock that work; if we
want laziness again we would want to reimplement it in C++ anyway.
2021-03-09 21:40:46 -05:00
Skye Wanderman-Milne
f06bb9a7f4 Update jaxlib version etc. 2021-03-09 17:55:40 -08:00
jax authors
23099f6007 Merge pull request #5835 from google:awn-abstract-eval
PiperOrigin-RevId: 361923732
2021-03-09 16:27:37 -08:00
Qiao Zhang
9577860169 Add jtu.count_device_put for tests to count device_put. 2021-03-09 14:45:01 -08:00
Matthew Johnson
2b9ffb1fb3 make axis_index bind respect dynamic traces 2021-03-09 13:51:12 -08:00
James Bradbury
47f17f0480 add (broken) test 2021-03-09 13:51:12 -08:00
James Bradbury
a8b8246554 add some todos 2021-03-09 13:51:09 -08:00
Roy Frostig
1f17fa6aa9 avoid modifying global pytype mappings in named-shape jaxpr test 2021-03-09 13:48:26 -08:00
Roy Frostig
e779ed8299 simplify standard named_shape_rule
Co-authored-by: Matthew Johnson <mattjj@google.com>
2021-03-09 13:48:26 -08:00
James Bradbury
c622422dad [avals with names] Propagate presence of name (mapped) vs absence (replicated) in abstract eval based on existing batching rules 2021-03-09 13:48:15 -08:00
James Bradbury
ba0f785a1f allow named axes on while_loop condition aval 2021-03-09 13:45:17 -08:00
Jake VanderPlas
dbdb189de1 jnp.piecewise: support scalar inputs 2021-03-09 13:25:38 -08:00
jax authors
6515b5f676 Merge pull request #5977 from apaszke:xmap-with-control-flow
PiperOrigin-RevId: 361854852
2021-03-09 11:22:18 -08:00
jax authors
f1ba3bcc9d Merge pull request #5990 from jakevdp:fix-power
PiperOrigin-RevId: 361852538
2021-03-09 11:12:39 -08:00
Skye Wanderman-Milne
55c75c8ca3 Update cuDNN versions to highest available.
The TF build depends on newer versions than we previously used.
2021-03-09 19:07:00 +00:00
jax authors
b0d14fd28f Merge pull request #5951 from apaszke:revive-all-to-all
PiperOrigin-RevId: 361841703
2021-03-09 10:29:12 -08:00
Jake VanderPlas
0c86c1fd11 jnp.power: fix overflow case for x1=0 2021-03-09 09:36:41 -08:00
jax authors
0b88b0ea9b Merge pull request #5856 from google:autodidax
PiperOrigin-RevId: 361722343
2021-03-08 20:50:15 -08:00