https://github.com/google/jax

sort by:
Revision Author Date Message Commit Date
ee2a92e pretty happy with this! 22 October 2019, 01:17:37 UTC
0cbad47 seems to work! 21 October 2019, 02:10:30 UTC
35f04c4 welp 15 October 2019, 03:37:36 UTC
5076a61 aaaaah 12 October 2019, 03:29:01 UTC
74eb0d8 wip 10 October 2019, 03:14:26 UTC
1de5688 wip optimizers 09 October 2019, 19:32:53 UTC
e8c9815 back to a primitive 09 October 2019, 19:32:53 UTC
1c2e635 more general/fewer primitives 09 October 2019, 19:32:53 UTC
c83245a [WIP] final-style unzip for key-split-based NNs 09 October 2019, 19:32:53 UTC
85cf558 experiments on stax Co-authored-by: Dougal Maclaurin <dougalm@google.com> 09 October 2019, 19:32:53 UTC
54807b4 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 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 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 add test for the case that failed 08 October 2019, 18:08:42 UTC
bc0e797 fix XLA metadata for primitives with many args 08 October 2019, 17:57:36 UTC
59ae24e Remove build_mac_jaxlib.yml for now to see if that's causing spurious workflows. 08 October 2019, 16:57:21 UTC
d3fa506 Bump jaxlib version to 0.1.31 and update WORKSPACE. 08 October 2019, 16:38:39 UTC
6c6e90d Merge pull request #1455 from gnecula/colabs Increase readthedocs/nbsphinx timeout 08 October 2019, 15:41:36 UTC
c6ac94f Merge pull request #1457 from shoyer/readthedocs-yml Add a .readthedocs.yml file 08 October 2019, 15:20:02 UTC
b7335be update to version 2 schema 07 October 2019, 23:57:00 UTC
cd29780 minor README formatting and grammar adjustment 07 October 2019, 23:55:17 UTC
f30fc7c 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 Merge pull request #1446 from joaogui1/selu [Done] Fixed selu activation 07 October 2019, 19:00:54 UTC
4145763 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 Merge pull request #1423 from gnecula/colabs Added How_JAX_primitives_work colab 07 October 2019, 12:59:37 UTC
505caaa fixed selu 05 October 2019, 01:39:35 UTC
a0bb2c0 Add a pylintrc to make it easier to use linter (#1442) 05 October 2019, 01:19:31 UTC
8891119 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 remove source lines (caching makes them ~useless) 04 October 2019, 18:59:06 UTC
c4e4199 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 Update docs/notebooks/How_JAX_primitives_work.ipynb Co-Authored-By: James Bradbury <jekbradbury@google.com> 04 October 2019, 09:58:07 UTC
ef886c7 Update docs/notebooks/How_JAX_primitives_work.ipynb Co-Authored-By: James Bradbury <jekbradbury@google.com> 04 October 2019, 09:56:06 UTC
f205213 Fix typo in initial_step_size d2 calculation (#1428) 04 October 2019, 06:43:34 UTC
f04f06f Merge pull request #1427 from trevorcai/doc Move epsilon term in optimizers.rmsprop into sqrt 04 October 2019, 01:12:02 UTC
b1120e4 Merge pull request #1436 from google/jb/ortho-init-type Fix initializers and add tests 04 October 2019, 01:11:08 UTC
59343b1 provide lax primitive + src line as XLA debuginfo 04 October 2019, 00:56:25 UTC
5201c36 Typo: Gradeints -> Gradients 04 October 2019, 00:48:33 UTC
e8bf56d 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 Update README to point to jaxlib-0.1.30 03 October 2019, 21:21:44 UTC
db3b0dd fix missing delimiter in notebook source 03 October 2019, 21:02:45 UTC
d13599e Merge pull request #1425 from gnecula/colab_testing Added support for testing Colab notebooks. 03 October 2019, 19:23:12 UTC
aa35f8b more initializer aliases 03 October 2019, 19:06:07 UTC
23f06b4 add initializer tests 03 October 2019, 19:01:21 UTC
4210192 fix two typos in orthogonal init 03 October 2019, 19:01:13 UTC
0fdf377 Switched testing to using sphinx (part of the documentation generation) 03 October 2019, 18:54:28 UTC
57e39ce 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 Bump jaxlib version to 0.1.30 and update WORKSPACE. 03 October 2019, 17:36:29 UTC
22e3251 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 Fix the title to show nicely in sphinx 03 October 2019, 09:59:51 UTC
75c2236 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 unbreak debug_nans on jit 02 October 2019, 23:35:42 UTC
dd1e132 test that debug_nans actually finds NaNs 02 October 2019, 23:10:27 UTC
d174fd6 fix config.update when absl is enabled 02 October 2019, 23:08:05 UTC
efc3a2c address comments 02 October 2019, 22:55:09 UTC
be4e156 use setUp and tearDown 02 October 2019, 22:51:15 UTC
459bd96 add basic tests 02 October 2019, 22:02:15 UTC
ee3fe5f fix debug_nans for multiple-result primitives 02 October 2019, 21:15:11 UTC
c18d765 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 Merge pull request #1424 from necula01/documentation1 Fixed links to Colabs in docs/notebooks 02 October 2019, 14:44:16 UTC
6744389 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 Added How_JAX_primitives_work colab 02 October 2019, 12:42:01 UTC
db1c04b Merge pull request #1420 from trevorcai/tn_minval Use finfo(dtype).tiny as uniform minval 02 October 2019, 07:20:18 UTC
a44c1ca 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 Fixup colab notebook URLs (#1417) * Fixup colab notebook URLs * wording 30 September 2019, 18:29:27 UTC
29d18fd wording 30 September 2019, 18:17:59 UTC
9805447 wording 30 September 2019, 18:15:37 UTC
08254fe Fixup colab notebook URLs 30 September 2019, 18:12:01 UTC
eb0137b 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 Notebooks on RTD (#1121) 30 September 2019, 18:00:02 UTC
064014b Merge pull request #1374 from google/jb/abs-jvp Improve numerics of abs jvp (and softplus) 29 September 2019, 01:43:25 UTC
90fdf64 Merge pull request #1415 from hawkinsp/master Update README to point to jaxlib 0.1.29. 28 September 2019, 22:43:34 UTC
7a96e50 Update README to point to jaxlib 0.1.29. 28 September 2019, 22:42:52 UTC
e4c4414 Merge pull request #1413 from hawkinsp/master Update Jaxlib version to 0.1.29. 28 September 2019, 19:17:20 UTC
1428c11 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 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 Add jax.local_devices() and jax.host_ids(). 27 September 2019, 23:36:45 UTC
9cf7a06 Remove xla_bridge.shape_of and xla_bridge._JaxComputationBuilder.Parameter, which are unused. 27 September 2019, 22:10:19 UTC
9d4a158 tweak comment in lax.root 27 September 2019, 19:47:39 UTC
c33e8cb Merge pull request #1339 from shoyer/solvers lax.root, a primitive for differentiable root finding 27 September 2019, 19:17:47 UTC
03793a8 Fix pmap error message. 27 September 2019, 18:50:21 UTC
840effc nanmean + nanprod only work on floating-point types. 27 September 2019, 18:27:24 UTC
d6a9202 test_util._cast_to_shape fix 27 September 2019, 18:18:36 UTC
b82673d add check_dtypes 27 September 2019, 16:11:18 UTC
01f8187 [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 nanmean fix 27 September 2019, 00:10:49 UTC
796d369 Remove licenses() rule comment in BUILD files. Internal tooling doesn't like it. 26 September 2019, 21:54:07 UTC
dc2ee0d 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 Merge pull request #1404 from google/remove-extra-newline remove extra newline to help copybara 26 September 2019, 18:57:27 UTC
c04f566 remove extra newline to help copybara 26 September 2019, 18:56:22 UTC
2206be2 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 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 Googly indentation 26 September 2019, 11:39:35 UTC
806ebd4 Fix and test for fori_loop op-by-op caching 26 September 2019, 11:34:57 UTC
b24d6ca Ensure LU decomposition cache hits in op-by-op 26 September 2019, 09:25:51 UTC
41d8a72 whitespace 25 September 2019, 23:16:01 UTC
4b53ea2 slightly cleaner 25 September 2019, 23:13:01 UTC
760d696 error messages 25 September 2019, 23:04:15 UTC
854aa28 notation 25 September 2019, 23:02:25 UTC
8b1a60e Revert hack from lax_linalg.py 25 September 2019, 22:51:52 UTC
6fda7be revert structure of lax.root 25 September 2019, 22:48:49 UTC
back to top