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

13184 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Yash Katariya
9e4114f0f1 Move array.py and sharding.py from experimental/ to _src/. 
PiperOrigin-RevId: 477201711
jax-v0.3.19 		2022-09-27 10:06:52 -07:00
jax authors
0e116888ea Merge pull request #12382 from jakevdp:reduction-dtype 
PiperOrigin-RevId: 477179725
 2022-09-27 08:38:46 -07:00
jax authors
1bcf8d646d Merge pull request #12497 from mattjj:djax-dag-fix1 
PiperOrigin-RevId: 477038279
 2022-09-26 18:14:56 -07:00
jax authors
e42247bffb Merge pull request #12524 from sharadmv:lax-import-fix 
PiperOrigin-RevId: 477038211
 2022-09-26 18:08:45 -07:00
Yash Katariya
389a2e570d Add a backwards compat path for op_sharding.clone() because it doesn't exist with the latest jaxlib on pypi 
PiperOrigin-RevId: 477034758
 2022-09-26 17:50:19 -07:00
Matthew Johnson
1e7ca8f77a fix bug in djax type signature inference logic 
Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com>
 2022-09-26 17:48:25 -07:00
Sharad Vikram
1d895b2c85 Fix lax imports 2022-09-26 17:32:44 -07:00
Yash Katariya
cbf34cb609 Rename the concrete class Array to ArrayImpl 
PiperOrigin-RevId: 477017236
 2022-09-26 16:18:30 -07:00
Tianjian Lu
71bcabe499 [sparse] Add BCSR format template. 
PiperOrigin-RevId: 477013899
 2022-09-26 16:02:16 -07:00
jax authors
82636b0bcd Merge pull request #12523 from jakevdp:fix-build 
PiperOrigin-RevId: 477005157
 2022-09-26 15:23:02 -07:00
Jake VanderPlas
6cae54f82d Fix bazel build alias 2022-09-26 15:13:12 -07:00
Peter Hawkins
d63a9442bb Change jax_jit_test to be a jax_test() under Bazel that works across backends. 
Make it pass under TPU if x64 types are enabled.

PiperOrigin-RevId: 476994286
 2022-09-26 14:38:35 -07:00
Jake VanderPlas
265b39d23f Add pytype_srcs to main jax BUILD rule 
PiperOrigin-RevId: 476989241
 2022-09-26 14:18:13 -07:00
jax authors
ddd8581f38 Merge pull request #12480 from google:bug-template-gpu-smi 
PiperOrigin-RevId: 476979981
 2022-09-26 13:41:31 -07:00
Jake VanderPlas
1860f6d839 [x64] add promote_integers argument to jnp.prod & jnp.sum 2022-09-26 13:31:43 -07:00
jax authors
69d1a2c063 Merge pull request #12517 from skye:update-pypi 
PiperOrigin-RevId: 476969287
 2022-09-26 13:00:31 -07:00
Yash Katariya
b2b60d943e Add make_array_from_single_device_arrays to prepare to rename of the concrete Array to ArrayImpl. 
PiperOrigin-RevId: 476965287
 2022-09-26 12:43:59 -07:00
Skye Wanderman-Milne
3c0d280bc0 Update version and changelog for jax 0.3.18 release 2022-09-26 12:43:39 -07:00
Roy Frostig
2a7b3197e0 add nvidia-smi question to bug template 2022-09-26 11:06:29 -07:00
jax authors
e034432872 Merge pull request #12513 from inoryy:patch-4 
PiperOrigin-RevId: 476923412
jax-v0.3.18 		2022-09-26 10:04:14 -07:00
jax authors
7962b01f5d Merge pull request #12485 from LenaMartens:checkify-lower 
PiperOrigin-RevId: 476922387
 2022-09-26 09:53:40 -07:00
lenamartens
27e3981d52 lowerable errors behind a config flag. 2022-09-26 17:34:27 +01:00
Roman Ring
8bcf358fde
Remove unused _remat_static_argnums import. 2022-09-26 17:14:09 +01:00
lenamartens
78ecc1442c Lowerable checks!! 2022-09-26 16:54:18 +01:00
jax authors
28672cca0e Merge pull request #12496 from mattjj:improve-leak-checker-2 
PiperOrigin-RevId: 476907407
 2022-09-26 08:50:13 -07:00
jax authors
9c66569514 Merge pull request #12468 from LenaMartens:checkify-but-better 
PiperOrigin-RevId: 476901601
 2022-09-26 08:23:02 -07:00
jax authors
2df61b1aa1 Merge pull request #12421 from jakevdp:jax-array 
PiperOrigin-RevId: 476898184
 2022-09-26 08:07:11 -07:00
Jake VanderPlas
0cb233eec9 Add initial jax.Array base class for instance checks & annotation 2022-09-26 07:48:43 -07:00
jax authors
2180710c2a Merge pull request #12511 from hawkinsp:release 
PiperOrigin-RevId: 476889960
jax-v0.3.18-rc 		2022-09-26 07:24:44 -07:00
Peter Hawkins
bcd36d8eb2 Jax and jaxlib 0.3.18 release candidate. 2022-09-26 14:10:57 +00:00
jax authors
53de057748 Merge pull request #12510 from hawkinsp:context 
PiperOrigin-RevId: 476884674
 2022-09-26 06:58:46 -07:00
Peter Hawkins
f4bc663c31 Wrap multiprocess test popen() uses in a context manager. 
Ensures that resources from popen() are cleaned up.
 2022-09-26 13:48:56 +00:00
jax authors
ec15e83018 - Wraps calls to lax.xeinsum and _einsum in a named call with their 'spec', the string specifying the computation. Makes xprof traces more interpretable. 
PiperOrigin-RevId: 476796185
 2022-09-25 20:54:17 -07:00
Yash Katariya
7c85ca38f4 Only look at hlo_modules for output sharding if there is more than 1 device because if there is only 1 device, the spmd partitioner won't run. 
PiperOrigin-RevId: 476497929
 2022-09-23 17:31:33 -07:00
Peter Hawkins
8ee7129874 Fix jnp.unwrap() test failures on GPU. 
A recent XLA change allows XLA to use excess precision on GPU, which caused CompileAndCheck to report noticeable numerical changes for bfloat16.

In passing, also enable the comparison against NumPy test for bfloat16 by using a wrapper function.

PiperOrigin-RevId: 476494989
 2022-09-23 17:11:51 -07:00
jax authors
d2fcfb6b83 Merge pull request #12407 from hirwa-nshuti:docs-fix 
PiperOrigin-RevId: 476467728
 2022-09-23 14:51:11 -07:00
Matthew Johnson
03abcc7c5c fix typo in test 2022-09-23 14:43:24 -07:00
jax authors
e76aa77895 Merge pull request #12437 from sudhakarsingh27:add_multi_host_pjit_tests 
PiperOrigin-RevId: 476451469
 2022-09-23 13:38:59 -07:00
Yash Katariya
1fa0dda760 Return single device Arrays from .device_buffer and .device_buffers. 
PiperOrigin-RevId: 476449591
 2022-09-23 13:30:26 -07:00
jax authors
43bbce0cc6 Merge pull request #12486 from hawkinsp:debugging 
PiperOrigin-RevId: 476445041
 2022-09-23 13:09:26 -07:00
jax authors
737327a42d Merge pull request #12490 from mattjj:improve-leak-checker 
PiperOrigin-RevId: 476442352
 2022-09-23 12:58:03 -07:00
Matthew Johnson
b6ef90ffdd fix leak checker internal error 
The issue was that partial_eval.py's _memoize, used in custom_jvp, was made
into an identity function by enabling config.jax_check_tracer_leaks (from
references to the main trace (needed for the jvp_jaxpr thunk) and hence trigger
the leak checker (which would see if any references to the main trace persisted
after finishing tracing of the user function).

But after #7345, the leak checker should only trigger when actual Tracers are
leaked. So disabling the memoization when jax_check_tracer_leaks is no longer
active shouldn't be necessary. (These PR numbers seem out of order! We're not
sure why.)

Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com>
 2022-09-23 12:33:45 -07:00
Sudhakar
4dd0d85139 add multihost pjit tests 2022-09-23 12:11:56 -07:00
Jake VanderPlas
a6b24b379c Add regression test for lax.rev simplification error 
PiperOrigin-RevId: 476430486
 2022-09-23 12:07:15 -07:00
Yash Katariya
ecb27a9b24 Update the _check_special code to not use xla_shape since its deprecated and does not work with Array. 
PiperOrigin-RevId: 476422732
 2022-09-23 11:40:32 -07:00
jax authors
d078f3f5fc Merge pull request #12478 from sharadmv:sharding-docs 
PiperOrigin-RevId: 476420315
 2022-09-23 11:31:37 -07:00
jax authors
e8865c8264 Merge pull request #12481 from kho:changelist/476272494 
PiperOrigin-RevId: 476411483
 2022-09-23 10:55:10 -07:00
Ke Wu
c823151771 Allow transpose axes to be negative to match (undocumented) NumPy behavior 2022-09-23 10:18:23 -07:00
Tres Popp
0c085471c7 Modify CorrCoef test to not rely on floating poing representation of 1/3 
The operation computed an average while using the dimension of size 3. This is then changed to multiplying by 1/3 with compilers, but 1/3 cannot be represented perfectly. That made this test case rely on a very precise result from an unrepresentable calculation.

PiperOrigin-RevId: 476391389
 2022-09-23 09:39:01 -07:00
Peter Hawkins
38fb8ed22f Fix copyright attribution for some newly added files. 
PiperOrigin-RevId: 476390902
 2022-09-23 09:32:47 -07:00