ee2a92e | James Bradbury | 22 October 2019, 01:17:37 UTC | pretty happy with this! | 22 October 2019, 01:17:37 UTC |
0cbad47 | James Bradbury | 21 October 2019, 02:10:30 UTC | seems to work! | 21 October 2019, 02:10:30 UTC |
35f04c4 | James Bradbury | 15 October 2019, 03:37:36 UTC | welp | 15 October 2019, 03:37:36 UTC |
5076a61 | James Bradbury | 12 October 2019, 03:29:01 UTC | aaaaah | 12 October 2019, 03:29:01 UTC |
74eb0d8 | James Bradbury | 10 October 2019, 03:14:26 UTC | wip | 10 October 2019, 03:14:26 UTC |
1de5688 | James Bradbury | 09 October 2019, 01:55:13 UTC | wip optimizers | 09 October 2019, 19:32:53 UTC |
e8c9815 | James Bradbury | 21 September 2019, 05:00:52 UTC | back to a primitive | 09 October 2019, 19:32:53 UTC |
1c2e635 | James Bradbury | 20 September 2019, 20:58:28 UTC | more general/fewer primitives | 09 October 2019, 19:32:53 UTC |
c83245a | James Bradbury | 13 September 2019, 04:33:21 UTC | [WIP] final-style unzip for key-split-based NNs | 09 October 2019, 19:32:53 UTC |
85cf558 | Matthew Johnson | 28 August 2019, 16:13:45 UTC | experiments on stax Co-authored-by: Dougal Maclaurin <dougalm@google.com> | 09 October 2019, 19:32:53 UTC |
54807b4 | Skye Wanderman-Milne | 08 October 2019, 21:28:14 UTC | Revert "Add a pylintrc to make it easier to use linter (#1442)" This reverts commit a0bb2c0ea452975be76e0ba2c6055f5be4439aa3. Temporarily reverting this to see if it's causing the github workflow failures. | 08 October 2019, 21:28:14 UTC |
2308d0f | Skye Wanderman-Milne | 08 October 2019, 21:20:40 UTC | Revert "Remove build_mac_jaxlib.yml for now to see if that's causing spurious workflows." This reverts commit 59ae24e87459429f18b0acb6626a1fb3465c41c8. | 08 October 2019, 21:20:40 UTC |
d34d9eb | James Bradbury | 08 October 2019, 20:30:07 UTC | Merge pull request #1462 from google/jb/fix-primitive-metadata fix XLA metadata for primitives with many args | 08 October 2019, 20:30:07 UTC |
245d29b | James Bradbury | 08 October 2019, 18:08:42 UTC | add test for the case that failed | 08 October 2019, 18:08:42 UTC |
bc0e797 | James Bradbury | 08 October 2019, 17:57:36 UTC | fix XLA metadata for primitives with many args | 08 October 2019, 17:57:36 UTC |
59ae24e | Skye Wanderman-Milne | 08 October 2019, 16:57:21 UTC | Remove build_mac_jaxlib.yml for now to see if that's causing spurious workflows. | 08 October 2019, 16:57:21 UTC |
d3fa506 | Skye Wanderman-Milne | 08 October 2019, 16:38:32 UTC | Bump jaxlib version to 0.1.31 and update WORKSPACE. | 08 October 2019, 16:38:39 UTC |
6c6e90d | George Necula | 08 October 2019, 15:41:36 UTC | Merge pull request #1455 from gnecula/colabs Increase readthedocs/nbsphinx timeout | 08 October 2019, 15:41:36 UTC |
c6ac94f | George Necula | 08 October 2019, 15:20:02 UTC | Merge pull request #1457 from shoyer/readthedocs-yml Add a .readthedocs.yml file | 08 October 2019, 15:20:02 UTC |
b7335be | Stephan Hoyer | 07 October 2019, 23:57:00 UTC | update to version 2 schema | 07 October 2019, 23:57:00 UTC |
cd29780 | Roy Frostig | 07 October 2019, 23:55:17 UTC | minor README formatting and grammar adjustment | 07 October 2019, 23:55:17 UTC |
f30fc7c | Stephan Hoyer | 07 October 2019, 23:09:23 UTC | Add a .readthedocs.yml file I think we need the leading dot for RTD to recognize it. We currently get warnings in our build log: "Configure your documentation builds! Adding a .readthedocs.yml file to your project is the recommended way to configure your documentation builds. You can declare dependencies, set up submodules, and many other great features." | 07 October 2019, 23:09:23 UTC |
2b46e0d | James Bradbury | 07 October 2019, 19:00:54 UTC | Merge pull request #1446 from joaogui1/selu [Done] Fixed selu activation | 07 October 2019, 19:00:54 UTC |
4145763 | George Necula | 07 October 2019, 16:07:55 UTC | Increase readthedocs/nbsphinx timeout It seems tht RTD is timing out when compiling the How_JAX_primitives_work noteboook. | 07 October 2019, 16:07:55 UTC |
a5063ee | George Necula | 07 October 2019, 12:59:37 UTC | Merge pull request #1423 from gnecula/colabs Added How_JAX_primitives_work colab | 07 October 2019, 12:59:37 UTC |
505caaa | joaogui1 | 05 October 2019, 01:39:35 UTC | fixed selu | 05 October 2019, 01:39:35 UTC |
a0bb2c0 | joao guilherme | 05 October 2019, 01:19:31 UTC | Add a pylintrc to make it easier to use linter (#1442) | 05 October 2019, 01:19:31 UTC |
8891119 | James Bradbury | 04 October 2019, 21:43:10 UTC | Merge pull request #1438 from google/jb/xla-debug-info thread through some debug metadata to XLA | 04 October 2019, 21:43:10 UTC |
4e6385f | James Bradbury | 04 October 2019, 18:59:06 UTC | remove source lines (caching makes them ~useless) | 04 October 2019, 18:59:06 UTC |
c4e4199 | George Necula | 04 October 2019, 15:38:36 UTC | Fix typo in notebook JSON (#1441) Perhaps this is the reason why readthedocs does not show the link to the notebook? | 04 October 2019, 15:38:36 UTC |
73144fa | George Necula | 04 October 2019, 09:58:07 UTC | Update docs/notebooks/How_JAX_primitives_work.ipynb Co-Authored-By: James Bradbury <jekbradbury@google.com> | 04 October 2019, 09:58:07 UTC |
ef886c7 | George Necula | 04 October 2019, 09:56:06 UTC | Update docs/notebooks/How_JAX_primitives_work.ipynb Co-Authored-By: James Bradbury <jekbradbury@google.com> | 04 October 2019, 09:56:06 UTC |
f205213 | Dan Weaver | 04 October 2019, 06:43:34 UTC | Fix typo in initial_step_size d2 calculation (#1428) | 04 October 2019, 06:43:34 UTC |
f04f06f | James Bradbury | 04 October 2019, 01:12:02 UTC | Merge pull request #1427 from trevorcai/doc Move epsilon term in optimizers.rmsprop into sqrt | 04 October 2019, 01:12:02 UTC |
b1120e4 | James Bradbury | 04 October 2019, 01:11:08 UTC | Merge pull request #1436 from google/jb/ortho-init-type Fix initializers and add tests | 04 October 2019, 01:11:08 UTC |
59343b1 | James Bradbury | 04 October 2019, 00:56:25 UTC | provide lax primitive + src line as XLA debuginfo | 04 October 2019, 00:56:25 UTC |
5201c36 | Fabian Pedregosa | 04 October 2019, 00:20:50 UTC | Typo: Gradeints -> Gradients | 04 October 2019, 00:48:33 UTC |
e8bf56d | Stephan Hoyer | 03 October 2019, 21:27:23 UTC | Refactor pytree munging inside lax.root. (#1414) * Refactor pytree munging inside lax.root. The primitive's implementation and JVP rules are now 100% pytree free. * Linearize multiple times in _root_jvp to avoid zeros * Revisions per review | 03 October 2019, 21:27:23 UTC |
7f98682 | Skye Wanderman-Milne | 03 October 2019, 21:21:28 UTC | Update README to point to jaxlib-0.1.30 | 03 October 2019, 21:21:44 UTC |
db3b0dd | Roy Frostig | 03 October 2019, 21:02:45 UTC | fix missing delimiter in notebook source | 03 October 2019, 21:02:45 UTC |
d13599e | Stephan Hoyer | 03 October 2019, 19:23:12 UTC | Merge pull request #1425 from gnecula/colab_testing Added support for testing Colab notebooks. | 03 October 2019, 19:23:12 UTC |
aa35f8b | James Bradbury | 03 October 2019, 19:06:07 UTC | more initializer aliases | 03 October 2019, 19:06:07 UTC |
23f06b4 | James Bradbury | 03 October 2019, 19:01:21 UTC | add initializer tests | 03 October 2019, 19:01:21 UTC |
4210192 | James Bradbury | 03 October 2019, 19:01:13 UTC | fix two typos in orthogonal init | 03 October 2019, 19:01:13 UTC |
0fdf377 | George Necula | 02 October 2019, 19:44:14 UTC | Switched testing to using sphinx (part of the documentation generation) | 03 October 2019, 18:54:28 UTC |
57e39ce | George Necula | 02 October 2019, 15:51:00 UTC | Added support for testing Colab notebooks. Testing is done by running "jupyter nbconvert --to notebook" and then parsing the resulting notebook to look for errors. One can declare expected errors, and the test will fail if those are missing. In the process of doig this, found and fixed a bug in the autodiff_cookbook notebook. | 03 October 2019, 18:07:49 UTC |
2a8575b | Skye Wanderman-Milne | 03 October 2019, 17:36:29 UTC | Bump jaxlib version to 0.1.30 and update WORKSPACE. | 03 October 2019, 17:36:29 UTC |
22e3251 | James Bradbury | 03 October 2019, 17:22:23 UTC | Merge pull request #1429 from google/jb/eig-debug-nans fix debug_nans for multiple-result primitives | 03 October 2019, 17:22:23 UTC |
7bc0f70 | George Necula | 03 October 2019, 09:59:51 UTC | Fix the title to show nicely in sphinx | 03 October 2019, 09:59:51 UTC |
75c2236 | George Necula | 03 October 2019, 09:20:04 UTC | Addressed comments for the Colab. * Cleaned up use of section levels * Renamed ma to multiply_add and sq_add to square_add * Other minor clarifications * Separated the Colabs into Tutorials and Advanced Tutorials | 03 October 2019, 09:20:04 UTC |
5103845 | James Bradbury | 02 October 2019, 23:35:42 UTC | unbreak debug_nans on jit | 02 October 2019, 23:35:42 UTC |
dd1e132 | James Bradbury | 02 October 2019, 23:10:27 UTC | test that debug_nans actually finds NaNs | 02 October 2019, 23:10:27 UTC |
d174fd6 | James Bradbury | 02 October 2019, 23:08:05 UTC | fix config.update when absl is enabled | 02 October 2019, 23:08:05 UTC |
efc3a2c | James Bradbury | 02 October 2019, 22:55:09 UTC | address comments | 02 October 2019, 22:55:09 UTC |
be4e156 | James Bradbury | 02 October 2019, 22:51:15 UTC | use setUp and tearDown | 02 October 2019, 22:51:15 UTC |
459bd96 | James Bradbury | 02 October 2019, 22:02:15 UTC | add basic tests | 02 October 2019, 22:02:15 UTC |
ee3fe5f | James Bradbury | 02 October 2019, 21:15:11 UTC | fix debug_nans for multiple-result primitives | 02 October 2019, 21:15:11 UTC |
c18d765 | Trevor Cai | 02 October 2019, 18:18:49 UTC | Move epsilon term in optimizers.rmsprop into sqrt The TensorFlow rmsprop code puts epsilon into the sqrt. This makes a large difference on e.g. Deep RL settings, which tend to use very large values of epsilon. | 02 October 2019, 18:18:49 UTC |
e359738 | Stephan Hoyer | 02 October 2019, 14:44:16 UTC | Merge pull request #1424 from necula01/documentation1 Fixed links to Colabs in docs/notebooks | 02 October 2019, 14:44:16 UTC |
6744389 | George Necula | 02 October 2019, 13:48:33 UTC | Fixed links to Colabs in docs/notebooks This is related to recent change https://github.com/google/jax/pull/1121. | 02 October 2019, 13:48:33 UTC |
454320e | George Necula | 02 October 2019, 12:41:28 UTC | Added How_JAX_primitives_work colab | 02 October 2019, 12:42:01 UTC |
db1c04b | James Bradbury | 02 October 2019, 07:20:18 UTC | Merge pull request #1420 from trevorcai/tn_minval Use finfo(dtype).tiny as uniform minval | 02 October 2019, 07:20:18 UTC |
a44c1ca | Trevor Cai | 01 October 2019, 21:28:31 UTC | Use finfo(dtype).tiny as uniform minval Otherwise, using the default clopen uniform for truncated_normal introduces a slight shift of the empirical mean. | 01 October 2019, 21:28:31 UTC |
675bfd5 | Stephan Hoyer | 30 September 2019, 18:29:27 UTC | Fixup colab notebook URLs (#1417) * Fixup colab notebook URLs * wording | 30 September 2019, 18:29:27 UTC |
29d18fd | Stephan Hoyer | 30 September 2019, 18:17:59 UTC | wording | 30 September 2019, 18:17:59 UTC |
9805447 | Stephan Hoyer | 30 September 2019, 18:15:37 UTC | wording | 30 September 2019, 18:15:37 UTC |
08254fe | Stephan Hoyer | 30 September 2019, 18:12:01 UTC | Fixup colab notebook URLs | 30 September 2019, 18:12:01 UTC |
eb0137b | Skye Wanderman-Milne | 28 September 2019, 00:37:44 UTC | Add optional `tuple_args` argument to xla_computation. This is useful when using JAX to create an HLO module that is compiled and executed elsewhere. Also fixes a bug in the `tuple_args` logic. | 30 September 2019, 18:10:32 UTC |
9bd7330 | Stephan Hoyer | 30 September 2019, 18:00:02 UTC | Notebooks on RTD (#1121) | 30 September 2019, 18:00:02 UTC |
064014b | James Bradbury | 29 September 2019, 01:43:25 UTC | Merge pull request #1374 from google/jb/abs-jvp Improve numerics of abs jvp (and softplus) | 29 September 2019, 01:43:25 UTC |
90fdf64 | Peter Hawkins | 28 September 2019, 22:43:34 UTC | Merge pull request #1415 from hawkinsp/master Update README to point to jaxlib 0.1.29. | 28 September 2019, 22:43:34 UTC |
7a96e50 | Peter Hawkins | 28 September 2019, 22:42:52 UTC | Update README to point to jaxlib 0.1.29. | 28 September 2019, 22:42:52 UTC |
e4c4414 | Peter Hawkins | 28 September 2019, 19:17:20 UTC | Merge pull request #1413 from hawkinsp/master Update Jaxlib version to 0.1.29. | 28 September 2019, 19:17:20 UTC |
1428c11 | Peter Hawkins | 28 September 2019, 19:05:41 UTC | Update Jaxlib version to 0.1.29. Bump XLA version. Enable C++14 mode since it is required by the new XLA version. | 28 September 2019, 19:11:09 UTC |
a1fcae3 | Peter Hawkins | 28 September 2019, 18:01:00 UTC | Merge pull request #1410 from hawkinsp/master Remove xla_bridge.shape_of and xla_bridge._JaxComputationBuilder.Para… | 28 September 2019, 18:01:00 UTC |
9ac5aa1 | Skye Wanderman-Milne | 27 September 2019, 21:38:16 UTC | Add jax.local_devices() and jax.host_ids(). | 27 September 2019, 23:36:45 UTC |
9cf7a06 | Peter Hawkins | 27 September 2019, 22:08:55 UTC | Remove xla_bridge.shape_of and xla_bridge._JaxComputationBuilder.Parameter, which are unused. | 27 September 2019, 22:10:19 UTC |
9d4a158 | Matthew Johnson | 27 September 2019, 19:47:39 UTC | tweak comment in lax.root | 27 September 2019, 19:47:39 UTC |
c33e8cb | Matthew Johnson | 27 September 2019, 19:17:47 UTC | Merge pull request #1339 from shoyer/solvers lax.root, a primitive for differentiable root finding | 27 September 2019, 19:17:47 UTC |
03793a8 | Skye Wanderman-Milne | 27 September 2019, 18:50:21 UTC | Fix pmap error message. | 27 September 2019, 18:50:21 UTC |
840effc | Skye Wanderman-Milne | 27 September 2019, 18:27:24 UTC | nanmean + nanprod only work on floating-point types. | 27 September 2019, 18:27:24 UTC |
d6a9202 | Skye Wanderman-Milne | 27 September 2019, 18:17:55 UTC | test_util._cast_to_shape fix | 27 September 2019, 18:18:36 UTC |
b82673d | James Bradbury | 27 September 2019, 16:11:18 UTC | add check_dtypes | 27 September 2019, 16:11:18 UTC |
01f8187 | Xiayun Sun | 27 September 2019, 16:00:38 UTC | [Issue1379] convert number predicates to boolean (#1395) * convert number predicates to boolean by checking whether they are equal to zero * one more test * try relative import for python 2.7 * try fixing circular import in python 2.7 take 2 * simplify test; add test for wrong type * more comprehensive tests; block collection type * use numpy to avoid recursive imports | 27 September 2019, 16:00:38 UTC |
226c9e9 | Skye Wanderman-Milne | 27 September 2019, 00:10:49 UTC | nanmean fix | 27 September 2019, 00:10:49 UTC |
796d369 | Skye Wanderman-Milne | 26 September 2019, 21:51:24 UTC | Remove licenses() rule comment in BUILD files. Internal tooling doesn't like it. | 26 September 2019, 21:54:07 UTC |
dc2ee0d | Skye Wanderman-Milne | 19 September 2019, 18:02:34 UTC | Add support for multihost pmaps. All participating hosts are assumed to be running the same pmap code. Conceptually, this can be considered a single pmap over an array sharded on its leading pmapped dimension across the hosts. Each host passes its input shard to its pmapped function call, which returns the corresponding output shard (i.e. an array of the same leading dimension size). However, any collective operations will be run across the entire "global" array. If the `devices` argument to pmap is None, the pmap is assumed to be running across all hosts visible to XLA (as returned by jax.host_count()). Each host can pass in an input array of leading dimension size equal to or less than the number of devices local to that host. Note that this doesn't change the current behavior for single-host platforms. If `devices` are specified, the participating hosts are dictated by the devices' host_ids, and each host must pass in an input array of leading dim size equal to the number of local participating devices. Implementation-wise, each host independently compiles the computation, which we assume yields the same executable on all hosts (follow-up work will add more error checking). The hosts must know the global axis size of the sharded array, e.g. to provide the correct replica count to XLA. This is equal to the length of `devices` if specified, but if not, pmap is recursively called (with `devices` specified) to use `psum` to compute the global axis size. | 26 September 2019, 21:44:16 UTC |
e7b09cf | Matthew Johnson | 26 September 2019, 18:57:27 UTC | Merge pull request #1404 from google/remove-extra-newline remove extra newline to help copybara | 26 September 2019, 18:57:27 UTC |
c04f566 | Matthew Johnson | 26 September 2019, 18:56:22 UTC | remove extra newline to help copybara | 26 September 2019, 18:56:22 UTC |
2206be2 | Skye Wanderman-Milne | 19 September 2019, 00:21:57 UTC | Put XLA computation arguments in a tuple if there are many arguments. This works around a TPU issue. This is only done for computations produced by jit or pmap, not primitive or nested computations. | 26 September 2019, 17:57:55 UTC |
762b602 | Matthew Johnson | 26 September 2019, 13:48:42 UTC | Merge pull request #1394 from j-towns/fix-scatter-caching Ensure all ops get cache hits on second op-by-op mode call | 26 September 2019, 13:48:42 UTC |
57d9fd6 | Jamie Townsend | 26 September 2019, 11:39:35 UTC | Googly indentation | 26 September 2019, 11:39:35 UTC |
806ebd4 | Jamie Townsend | 26 September 2019, 11:34:57 UTC | Fix and test for fori_loop op-by-op caching | 26 September 2019, 11:34:57 UTC |
b24d6ca | Jamie Townsend | 26 September 2019, 09:25:51 UTC | Ensure LU decomposition cache hits in op-by-op | 26 September 2019, 09:25:51 UTC |
41d8a72 | Stephan Hoyer | 25 September 2019, 23:16:01 UTC | whitespace | 25 September 2019, 23:16:01 UTC |
4b53ea2 | Stephan Hoyer | 25 September 2019, 23:13:01 UTC | slightly cleaner | 25 September 2019, 23:13:01 UTC |
760d696 | Stephan Hoyer | 25 September 2019, 23:04:15 UTC | error messages | 25 September 2019, 23:04:15 UTC |
854aa28 | Stephan Hoyer | 25 September 2019, 23:02:25 UTC | notation | 25 September 2019, 23:02:25 UTC |
8b1a60e | Stephan Hoyer | 25 September 2019, 22:51:52 UTC | Revert hack from lax_linalg.py | 25 September 2019, 22:51:52 UTC |
6fda7be | Stephan Hoyer | 25 September 2019, 22:48:49 UTC | revert structure of lax.root | 25 September 2019, 22:48:49 UTC |