mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
12,056 Commits 506 Branches 263 Tags
Commit Graph

12056 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Roy Frostig
621dfb9347 remove unused stages.Lowered._xla_computation 
PiperOrigin-RevId: 456412379
 2022-06-21 20:04:54 -07:00
Roy Frostig
6a0c05fa20 remove unused stages.Compiled._xla_executable 
PiperOrigin-RevId: 456412184
 2022-06-21 19:59:16 -07:00
Yash Katariya
e5031d15de Disable xla sharding propagation test for SE because XLA sharding propagation is not supported on SE which is activate when out_axis_resources is not specified in pjit. 
PiperOrigin-RevId: 456391444
 2022-06-21 17:40:31 -07:00
jax authors
15a3798424 Merge pull request #11147 from jakevdp:x64-scipy-signal-test 
PiperOrigin-RevId: 456389803
 2022-06-21 17:32:15 -07:00
jax authors
1d6cbc93dd Merge pull request #11192 from sharadmv:release 
PiperOrigin-RevId: 456355961
 2022-06-21 14:48:41 -07:00
Sharad Vikram
217d898124 Update TF version for jaxlib build 2022-06-21 14:34:22 -07:00
jax authors
c6e3f00eb3 Merge pull request #11189 from jakevdp:fix-jvp-doc 
PiperOrigin-RevId: 456341327
 2022-06-21 13:46:01 -07:00
jax authors
f85886ba65 Merge pull request #11188 from jakevdp:doc-fix 
PiperOrigin-RevId: 456339293
 2022-06-21 13:40:32 -07:00
jax authors
2dbf3e0795 Merge pull request #11187 from sharadmv:release 
PiperOrigin-RevId: 456338899
 2022-06-21 13:34:56 -07:00
Jake VanderPlas
abcfaec6e3 DOC: clarify variable names 2022-06-21 13:20:53 -07:00
Sharad Vikram
9bd1bd67e0 Update versions for jax/jaxlib release 2022-06-21 12:57:28 -07:00
Jake VanderPlas
ed152953ef DOC: fix output repr in thinking_in_jax 2022-06-21 12:54:50 -07:00
jax authors
cd11aeca83 Merge pull request #11138 from jakevdp:x64-scipy-optimize-test 
PiperOrigin-RevId: 456327057
 2022-06-21 12:47:34 -07:00
jax authors
da46e268b2 Merge pull request #11152 from jakevdp:x64-jet-test 
PiperOrigin-RevId: 456325791
 2022-06-21 12:42:23 -07:00
jax authors
5669be791c Merge pull request #11184 from jakevdp:update-myst-nb 
PiperOrigin-RevId: 456324987
 2022-06-21 12:36:33 -07:00
Yash Katariya
bf1b29362a Fix up some comments after pmap+Array CL. 
PiperOrigin-RevId: 456310489
 2022-06-21 11:32:46 -07:00
Jake VanderPlas
997c90b3bb [x64] mention flag value in strict type promotion error 2022-06-21 10:55:24 -07:00
jax authors
2c34097671 Merge pull request #11183 from tmi:docs-fix-profiler 
PiperOrigin-RevId: 456285950
 2022-06-21 10:01:36 -07:00
Jake VanderPlas
3a8f478b0a [x64] make scipy_optimize_test compatible with strict dtype promotion 2022-06-21 09:29:08 -07:00
Jake VanderPlas
03f2189f90 [x64] make jax.scipy.signal compatible with strict dtype promotion. 
Also a fair bit of cleanup & refactoring of related code.
 2022-06-21 09:28:46 -07:00
Jake VanderPlas
893179fcfc [x64] make jet_test compatible with strict dtype promotion 2022-06-21 09:28:24 -07:00
jax authors
d7ec244c0c Merge pull request #11160 from jakevdp:x64-xmap-test 
PiperOrigin-RevId: 456274863
 2022-06-21 09:13:38 -07:00
Jake VanderPlas
b5b78b1052 CI: build docs with most recent myst-nb version 2022-06-21 08:44:23 -07:00
vojta tuma
4e057c9f23 DOC: update pprof install instructions 2022-06-21 11:15:14 +02:00
jax authors
30365b8bcc Merge pull request #10852 from LenaMartens:changelist/442814308 
PiperOrigin-RevId: 456115702
 2022-06-20 13:10:12 -07:00
Lena Martens
9167f7248a Checkify: support discharging checks from control-flow through effects. 
Currently supports scan and while-loop.
 2022-06-20 18:28:03 +01:00
George Necula
4b03ebf4f5 Fix overflow of large prng computation 
Fixes: #11010
 2022-06-20 10:48:15 +02:00
Matthew Johnson
a001c52f87 [dynamic-shapes] basic jvp working, including with broadcast 2022-06-18 13:38:48 -07:00
jax authors
b50d77cc98 Merge pull request #10576 from mattjj:new-remat-landing 
PiperOrigin-RevId: 455831386
 2022-06-18 12:33:52 -07:00
Matthew Johnson
83a8dc4e7f [new-remat] add _scan_partial_eval_custom rule for new remat 
Also enable scan-of-remat tests which weren't passing before.

Co-authored-by: Sharad Vikram <sharadmv@google.com>
 2022-06-17 23:15:14 -07:00
jax authors
4f5115c288 Merge pull request #11161 from schmrlng:patch-2 
PiperOrigin-RevId: 455748673
 2022-06-17 20:45:48 -07:00
jax authors
57e3f24130 Merge pull request #11163 from jakevdp:x64-polynomial-test 
PiperOrigin-RevId: 455724347
 2022-06-17 17:11:41 -07:00
jax authors
9b7cad1db9 Merge pull request #11158 from jakevdp:x64-batching-test 
PiperOrigin-RevId: 455722926
 2022-06-17 17:01:29 -07:00
jax authors
812ff25796 Merge pull request #11150 from jakevdp:x64-scipy-sparse 
PiperOrigin-RevId: 455722856
 2022-06-17 17:01:17 -07:00
jax authors
76387ae961 Merge pull request #11162 from jakevdp:x64-sharded-jit 
PiperOrigin-RevId: 455722197
 2022-06-17 16:56:18 -07:00
Jake VanderPlas
15a19969de [x64] make polynomial_test compatible with strict dtype promotion 2022-06-17 16:52:22 -07:00
jax authors
5b3ed3036d Merge pull request #11159 from jakevdp:x64-pjit-test 
PiperOrigin-RevId: 455721480
 2022-06-17 16:51:00 -07:00
Jake VanderPlas
1b2743d63d [x64] make sharded_jit_test compatible with strict dtype promotion 2022-06-17 16:42:45 -07:00
Ed Schmerling
f2a85f36f5
Remove spurious single quote in profiler.py 
Before:
```
127.0.0.1 - - [17/Jun/2022 16:08:18] "GET /perfetto_trace.json.gz' HTTP/1.1" 404 -
```
After:
```
127.0.0.1 - - [17/Jun/2022 16:35:54] "GET /perfetto_trace.json.gz HTTP/1.1" 200 -
```
 2022-06-17 16:39:21 -07:00
Jake VanderPlas
c467b15f06 [x64] make xmap_test compatible with strict dtype promotion 2022-06-17 16:35:42 -07:00
Jake VanderPlas
51da9eb237 [x64] make pjit_test compatible with strict dtype promotion 2022-06-17 16:25:36 -07:00
jax authors
cc9ffbcd0f Merge pull request #11157 from jakevdp:x64-vmap-test 
PiperOrigin-RevId: 455717063
 2022-06-17 16:23:53 -07:00
Jake VanderPlas
9edbf9bb0d [x64] make batching_test compatible with strict dtype promotion 2022-06-17 16:08:54 -07:00
jax authors
5236140b9f Merge pull request #11155 from jakevdp:x64-ode-test 
PiperOrigin-RevId: 455714054
 2022-06-17 16:05:39 -07:00
Jake VanderPlas
7d21d6ef9d [x64] make vmap_test compatible with strict dtype promotion 2022-06-17 15:55:08 -07:00
jax authors
5318df67fc Merge pull request #11151 from mattjj:input-fwd-residual-optimization 
PiperOrigin-RevId: 455710192
 2022-06-17 15:43:44 -07:00
Matthew Johnson
72a67906bf optimize grad-of-jit not to pass input-residuals as intermediate-residuals 
Co-authored-by: Dougal Maclaurin <dougalm@google.com>
Co-authored-by: Peter Hawkins <phawkins@google.com>
 2022-06-17 15:24:28 -07:00
jax authors
5f849d3aaa Merge pull request #11116 from mattjj:djax-typecheck 
PiperOrigin-RevId: 455706708
 2022-06-17 15:23:19 -07:00
Jake VanderPlas
5e97011dc8 [x64] make ode_test pass with strict dtype promotion 2022-06-17 15:21:10 -07:00
Matthew Johnson
f680269a4f [dynamic-shapes] initial support for dynamic shape typechecks 
Co-authored-by: Dougal Maclaurin <dougalm@google.com>
 2022-06-17 14:57:19 -07:00