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

69 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Dougal Maclaurin
307d195577 Wrapped static args to jit to be hashed on id. This is conservative but simple and cheap. 2018-12-05 15:55:01 -05:00
Dougal Maclaurin
5bb8f87e22 First pass at a quickcheck-style property-based tester to test invariants of jit/grad/vmap etc on random functions 2018-12-05 15:55:01 -05:00
Matthew Johnson
2597300c7f source sync 
PiperOrigin-RevId: 224157599
 2018-12-05 09:36:16 -08:00
Matthew Johnson
3b049e3853 [XLA] separate out an Execute from ExecutePerReplica 
[JAX] reduce the creation of XLA Shape protos on every call (which is slow)

PiperOrigin-RevId: 223915944
 2018-12-05 09:36:12 -08:00
Peter Hawkins
77db9bd556 [JAX] Update XLA in JAX workspace to include optimized computation launch implementation. 
PiperOrigin-RevId: 223863753
 2018-12-05 09:36:08 -08:00
Roy Frostig
f5b051b431 Double gpu test shards 
PiperOrigin-RevId: 223647959
 2018-12-02 11:50:46 -08:00
Roy Frostig
20878c76f4 source sync 
PiperOrigin-RevId: 223530503
 2018-12-02 11:50:39 -08:00
Dougal Maclaurin
2df36f7510 Made a shim to handle configuration without having absl parse command-line flags. 
PiperOrigin-RevId: 223391288
 2018-11-29 13:44:54 -08:00
Matthew Johnson
1d2aaad6fe fix handling of symbolic zeros for a few special primitives 
PiperOrigin-RevId: 223264329
 2018-11-29 12:06:17 -08:00
Matthew Johnson
aefd8f9eb5 add contributing.md 
PiperOrigin-RevId: 223085949
 2018-11-27 17:18:37 -08:00
Roy Frostig
14844224fc source sync 
PiperOrigin-RevId: 223081582
 2018-11-27 16:51:28 -08:00
Matthew Johnson
0ea98501aa source sync 
PiperOrigin-RevId: 223080639
 2018-11-27 16:51:26 -08:00
Peter Hawkins
93ac03ea08 [JAX] Update TensorFlow release to 0b6ed4887a to pick up fixes to the XLA Mac OS X build. 
PiperOrigin-RevId: 223048563
 2018-11-27 16:51:24 -08:00
Dougal Maclaurin
ca2634ea5d source sync 
PiperOrigin-RevId: 222923229
 2018-11-27 16:51:22 -08:00
Peter Hawkins
599ea38175 [JAX] Explicitly use /bin/bash in install_xla_in_source_tree.sh. 
It turns out the script uses some bash-isms not supported by the /bin/sh (dash) shell on a Debian machine.

PiperOrigin-RevId: 222884583
 2018-11-27 16:51:19 -08:00
Peter Hawkins
f3513a7bfb [JAX] Rewrite OSS build script. 
Significant changes:
* Mac OS X support.
* build script is in Python, not shell.
* build configuration is passed via flags, not environment variables.
* build script configures TF itself, and does not require explicitly checking out the TF git repository and running its configure script. Changes the TF dependency in the Bazel workspace to be an http_archive(), rather than a local checkout of TF.
* rather than trying to guess the path for Bazel-generated XLA artifacts, use a sh_binary() to perform installation of the built artifacts in to the JAX source tree. Bazel's runfiles mechanism is the supported route to find build artifacts.
* downloads Bazel in Python and checks its SHA256 before running it, rather than running an untrusted binary from the internet.
* intentionally does not delete the Bazel cache or Bazel after building.

Example of new build interaction:

Building without CUDA on Mac or Linux:
$ cd jax
$ python3 build.py   (or python2 build.py if you want a Python 2 build)

     _   _    __  __
    | | / \   \ \/ /
 _  | |/ _ \   \  /
| |_| / ___ \  /  \
 \___/_/   \_\/_/\_\

Starting local Bazel server and connecting to it...
Bazel binary path: /Users/xyz/bin/bazel
Python binary path: /Library/Frameworks/Python.framework/Versions/3.7/bin/python3
CUDA enabled: no

Building XLA and installing it in the JAX source tree...
...

Example of building with CUDA enabled on Linux:
$ python3 build.py --enable_cuda --cudnn_path=/usr/lib/x86_64-linux-gnu/
... as before, except ...
CUDA enabled: yes
CUDA toolkit path: /usr/local/cuda
CUDNN library path: /usr/lib/x86_64-linux-gnu/
...

PiperOrigin-RevId: 222868835
 2018-11-27 16:51:17 -08:00
Sam Schoenholz
326773808b source sync 
PiperOrigin-RevId: 222867729
 2018-11-27 16:51:15 -08:00
Matthew Johnson
c293f7c875 minor: add @jit to threefry hash function in random.py 
PiperOrigin-RevId: 222841601
 2018-11-27 16:51:13 -08:00
Matthew Johnson
d318c68827 source sync 
PiperOrigin-RevId: 222822439
 2018-11-27 16:51:11 -08:00
Peter Hawkins
6001fd219d [JAX] Python 3 fix: xla_client.initialize_platform_name() accepts a string, not bytes, so we shouldn't pass it bytes. (The previous Python 3 change was a bit overzealous and made both changes; apparently this path isn't tested by the tests because they pass an explicit --jax_xla_backend flag..) 
PiperOrigin-RevId: 222810811
 2018-11-27 16:51:08 -08:00
Roy Frostig
642046e2c0 Bridge to the XRT backend. 
PiperOrigin-RevId: 222763810
 2018-11-27 16:51:06 -08:00
Matthew Johnson
dbf3c606f5 source sync 
PiperOrigin-RevId: 222500675
 2018-11-27 16:51:00 -08:00
Roy Frostig
7a313bf622 source sync fixups 2018-11-21 20:32:33 -08:00
Roy Frostig
f1bb77dafb source sync 
PiperOrigin-RevId: 222489712
 2018-11-21 20:32:33 -08:00
Matthew Johnson
9101d66b4e source sync 
PiperOrigin-RevId: 222487460
 2018-11-21 20:32:33 -08:00
Matthew Johnson
98dd8775d7 source sync 
PiperOrigin-RevId: 222484959
 2018-11-21 20:32:33 -08:00
Matthew Johnson
8317cc3618 source sync 
PiperOrigin-RevId: 222484671
 2018-11-21 20:32:33 -08:00
Matthew Johnson
3b3490f406 source sync 
PiperOrigin-RevId: 222483357
 2018-11-21 20:32:33 -08:00
Matthew Johnson
ab53373665 source sync 
PiperOrigin-RevId: 222470141
 2018-11-21 20:32:33 -08:00
Matthew Johnson
2ae9a2bc35 source sync 
PiperOrigin-RevId: 222461242
 2018-11-21 20:32:16 -08:00
Peter Hawkins
6361b784a8 source sync 
PiperOrigin-RevId: 222456068
 2018-11-21 20:22:56 -08:00
Peter Hawkins
5e60639bc5 source sync 
PiperOrigin-RevId: 222452709
 2018-11-21 20:22:54 -08:00
Peter Hawkins
e180f08113 source sync 
PiperOrigin-RevId: 222451919
 2018-11-21 20:22:51 -08:00
Peter Hawkins
fe4edf2839 source sync 
PiperOrigin-RevId: 222449830
 2018-11-21 20:22:49 -08:00
Matthew Johnson
323be694a7 source sync 
PiperOrigin-RevId: 222448341
 2018-11-21 20:22:47 -08:00
Matthew Johnson
e5b76f4fde source sync 
PiperOrigin-RevId: 222340967
 2018-11-21 20:22:43 -08:00
Peter Hawkins
065bb0baa2 source sync 
PiperOrigin-RevId: 222291726
 2018-11-21 20:22:41 -08:00
Roy Frostig
a954389d06 source sync 
PiperOrigin-RevId: 222189611
 2018-11-21 20:22:37 -08:00
Matthew Johnson
7f546b8c02 source sync 
PiperOrigin-RevId: 222175432
 2018-11-21 20:22:35 -08:00
Matthew Johnson
25fb9b421d source sync 
PiperOrigin-RevId: 222170151
 2018-11-21 20:22:33 -08:00
Roy Frostig
a3619ca89d source sync 
PiperOrigin-RevId: 222153576
 2018-11-21 20:22:30 -08:00
Matthew Johnson
51fc713089 remove absl from examples, fix import statements 2018-11-21 12:10:31 -08:00
Roy Frostig
7cf6babc78 Add two-pager link to README. 2018-11-21 08:57:25 -08:00
Matthew Johnson
b58ccb430a
update readme with 'not a Google product' 2018-11-21 06:04:06 -08:00
Matthew Johnson
76346a87b5
add placeholder readme.md 2018-11-21 06:00:31 -08:00
Matthew Johnson
377322d3d4 remove stray file 2018-11-19 21:58:15 -08:00
Matthew Johnson
5e9fcbb6e8 source sync 2018-11-19 21:31:39 -08:00
Matthew Johnson
b2b1e8d70c source sync 2018-11-19 21:24:30 -08:00
Matthew Johnson
50038c07c8 fix build file issues 2018-11-19 20:18:31 -08:00
Roy Frostig
fe11b19e46 source sync 2018-11-19 15:44:16 -08:00