mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 00:56:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
3,569 Commits 508 Branches 263 Tags
Commit Graph

29 Commits

Author SHA1 Message Date
Peter Hawkins
991324f8df
Increase minimum jaxlib version to 0.1.38. (#2120) 2020-01-29 14:16:58 -05:00
Peter Hawkins
681ba37f7e
Drop fastcache dependency, which isn't necessary on Python 3. (#1995) 
Drop protobuf and six dependencies from travis configuration.
 2020-01-14 10:08:23 -05:00
Peter Hawkins
dcc882cf6b
Drop Python 2 support from JAX. (#1962) 
Remove six dependency.
 2020-01-08 13:17:55 -05:00
Matthew Johnson
2a394ce31b move multi-device test into its own file 2019-12-18 14:40:20 -08:00
Matthew Johnson
00e621df75 add multi_device_test in travis 2019-12-18 14:40:20 -08:00
Matthew Johnson
286ec51f61 make op-by-op computation follow arg placement 2019-12-18 14:40:20 -08:00
Skye Wanderman-Milne
b459b6098a
Make pmap properly replicate closed-over constants. (#1847) 
With this change, a value `x` can be replicated `nrep` times as follows:

```python
pmap(lambda _: x)(np.arange(nrep))
```

This will broadcast `x` into a ShardedDeviceArray suitable for passing into another pmap with the same input shape.

If `x` will be passed into a pmap with `devices` or a nested pmap, the replication pmap(s) should follow that structure. For example:

```python
x = pmap(pmap(lambda _: x))(np.ones(2, 4))

pmap(pmap(lambda i: i**2 + x))(np.ones(2, 4))
```
 2019-12-17 16:22:55 -08:00
Peter Hawkins
e87d9718c3
Support IntEnum values as arguments to JAX functions. (#1840) 
* Support IntEnum values as arguments to JAX functions.

When abstractifying a Python value, search the method-resolution order (MRO) of the type rather than only looking at the value's own type. IntEnum instances are subclasses of int, so this allows us to correctly handle them as integers, much as NumPy itself does.
 2019-12-11 12:27:11 -05:00
Peter Hawkins
7a9f1f3f1c
Pin the minimum jaxlib version in travis.yml. (#1767) 2019-11-26 10:24:06 -05:00
George Necula
0ffcd769ef
Add sklearn to Travis, for documentation building. (#1547) 
* Add sklearn to Travis, for documentation building.
* Add score_matching to auto-built notebooks
 2019-10-21 23:24:16 +02:00
Matthew Johnson
0eea988502 try following jupyterlab/jupyter-renderers#212 to fix travis 2019-10-15 23:07:31 +00:00
George Necula
858a411982
Add a separate build matrix entry for documentation testing. (#1495) 
* Add a separate build matrix entry for documentation testing.

This way we parallelize the unit tests with the documentation tests.
 2019-10-13 09:34:51 +02:00
George Necula
0fdf377343 Switched testing to using sphinx (part of the documentation generation) 2019-10-03 20:54:28 +02:00
George Necula
57e39ce938 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.
 2019-10-03 20:07:49 +02:00
Peter Hawkins
08013954a4 Use fastcache for LRU caches in JAX. 
fastcache is both a faster cache implementation and is also thread-safe.
 2019-07-22 17:24:10 -04:00
Matthew Johnson
9790890b7d reduce travis sampled test cases from 100 to 25 
This change is just to speed up travis, which has ballooned from ~10
minutes to ~40 minutes.
 2019-07-02 12:38:16 -07:00
Matthew Johnson
d622e78e82 fix broken travis path from previous commit 2019-06-07 16:00:11 -07:00
Skye Wanderman-Milne
04b0a911c3 pip wheels need to include the versionso use version '0.1.18' instead of 'latest'. 2019-06-07 15:13:49 -07:00
Matthew Johnson
6fe6cb0dbe try fixing travis by using only one pytest job 2019-05-24 19:02:40 -07:00
Matthew Johnson
2802c8f514 try pinning travis to jaxlib==0.1.15 
Currently Travis CI builds are failing in a way we can't reproduce
locally.
 2019-05-24 10:44:33 -07:00
Matthew Johnson
6c7dde59f5 move isclose test behind numpy version check 2019-04-30 09:19:48 -07:00
Peter Hawkins
a355da5983 Change Travis configuration to use pytest-xdist. 
Since the Travis VM is dual core, run two tests in parallel.
 2019-04-12 11:27:32 -04:00
Matthew Johnson
77f0fca7fa fix np.polyval complex128 bug in x64=True mode 
Also enable JAX_ENABLE_X64=True tests on Travis to avoid future such
issues. (Internal tests catch things like this, but we don't run those
automatically.)
 2019-01-05 09:47:06 -08:00
Matthew Johnson
13b8e21a1c squash conv grad bug introduced in 0d64aea 
(loudly errored, didn't produce silently incorrect results!)
 2018-12-14 18:40:50 -08:00
Matthew Johnson
6de5c8a698 add test-running instructions (fixes #67) 2018-12-14 16:48:08 -08:00
Matthew Johnson
9b645364c9 tests depend on scipy, tweaks 2018-12-13 16:40:03 -08:00
Matthew Johnson
5d7f356bf5 use 'dist: xenial' in travis.yml to get updated toolchain 2018-12-13 16:21:18 -08:00
Matthew Johnson
b7728d4239 add jaxlib install to travis file 2018-12-13 13:41:49 -08:00
Matthew Johnson
32f89e6c5a add initial travis ci file 2018-12-13 13:39:34 -08:00