mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
11,554 Commits 504 Branches 263 Tags
Commit Graph

11554 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Yash Katariya
1381afc37f Update version after jax release 
PiperOrigin-RevId: 448822949
 2022-05-15 12:14:26 -07:00
jax authors
42421c81b4 Merge pull request #10688 from sharadmv:color-by-default 
PiperOrigin-RevId: 448723842
jax-v0.3.11 		2022-05-14 15:22:31 -07:00
Sharad Vikram
b8a523f977 Enable colors when we are using a terminal or IPython 2022-05-14 14:54:11 -07:00
jax authors
f26133cc1e Merge pull request #10584 from mattjj:remove-jnp-array-handling-raw-buffers 
PiperOrigin-RevId: 448720084
 2022-05-14 14:37:46 -07:00
jax authors
86899ee55f Merge pull request #10711 from mattjj:closed-call 
PiperOrigin-RevId: 448720052
 2022-05-14 14:32:39 -07:00
Matthew Johnson
05dda56019 add core.closed_call_p 2022-05-14 14:06:30 -07:00
Matthew Johnson
c0d6a04b76 remove jnp.array case for handling buffers w/ aval=None 
This functionality was added in #8134, but was superceded by later changes
which ensured that we never produce DeviceArrays with their 'aval' property set
to None (even when indexing ShardedDeviceArrays with integers, which used to be
a problem case).
 2022-05-14 08:21:54 -07:00
Hyojun Kim
ba0a2b3a98 Add profile_version to CompileOptions 
profile_version is added to CompileOptions to support XLA compilation profile.

PiperOrigin-RevId: 448599962
 2022-05-13 16:24:29 -07:00
Peter Hawkins
db73670ec3 Add support for padded arrays in QDWH algorithm. 
This change is in preparation for adding a jit-table QDWH-eig implementation.

PiperOrigin-RevId: 448571523
 2022-05-13 13:57:36 -07:00
jax authors
7c582ab970 Merge pull request #10680 from jakevdp:rm-future 
PiperOrigin-RevId: 448555184
 2022-05-13 12:38:05 -07:00
Jake VanderPlas
4d22d228b6 Fix version skew in api_test.py 
Issue was introduced in 8b073d482e

PiperOrigin-RevId: 448546982
 2022-05-13 11:58:45 -07:00
jax authors
bd659df4fe Merge pull request #10686 from mattjj:tweak-mlir-shape-tensor-helper 
PiperOrigin-RevId: 448523702
 2022-05-13 10:19:23 -07:00
Shuangchi He
8b073d482e PR #55768: Fix typos for occured, appearence, this, is, a, for, agressiveness, t… 
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/55768

Fix typos for occured, appearence, this, is, a, for, agressiveness, to, instrution, on.
Copybara import of the project:

--
531b97d4b242a5642a221349ca0bd3132d6539a2 by Yulv-git <yulvchi@qq.com>:

Fix typos for occured, appearence, this, is, a, for, agressiveness, to, instrution, on.

Merging this change closes #55768

PiperOrigin-RevId: 448444384
 2022-05-13 02:23:14 -07:00
Matthew Johnson
d719e5a55f tweak mlir shape_tensor helper, fewer MHLO ops 2022-05-12 17:02:08 -07:00
jax authors
2100c3b5be Merge pull request #10681 from JeppeKlitgaard:fix-int-typo 
PiperOrigin-RevId: 448313566
 2022-05-12 12:40:44 -07:00
jax authors
b9be65c107 Merge pull request #10679 from jakevdp:fix-dollarmath 
PiperOrigin-RevId: 448303782
 2022-05-12 11:58:00 -07:00
Jake VanderPlas
f92558c823 Remove unnecessary __future__ imports 2022-05-12 11:57:21 -07:00
Jeppe Klitgaard
e422441a65 fix: mypy warning 2022-05-12 19:45:41 +01:00
Jake VanderPlas
681240244c DOC: explicitly enable dollarmath extension 2022-05-12 10:37:34 -07:00
jax authors
41e1635ac1 Merge pull request #10664 from PhilipVinc:patch-1 
PiperOrigin-RevId: 448282864
 2022-05-12 10:30:11 -07:00
jax authors
0d88bdf00f Merge pull request #10642 from lgeiger:reverse-trivial-gather 
PiperOrigin-RevId: 448279969
 2022-05-12 10:18:10 -07:00
Peter Hawkins
7ba36fc178 Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh. 
Remove jax._src.lax.polar.

PiperOrigin-RevId: 448241206
 2022-05-12 07:20:52 -07:00
Peter Hawkins
880cfc9c79 Fix pseudocode in lax.switch() documentation to accept variadic operands. 
PiperOrigin-RevId: 448241093
 2022-05-12 07:15:46 -07:00
jax authors
8e4cf253dc Merge pull request #10616 from mattjj:partial-eval-jaxpr-custom-upgrade 
PiperOrigin-RevId: 448110206
 2022-05-11 16:27:58 -07:00
jax authors
6b2d1baf7a Merge pull request #10595 from mattjj:core-test-with-corecall 
PiperOrigin-RevId: 448109351
 2022-05-11 16:23:06 -07:00
Xin Zhou
06e4fbc48a [mhlo] Add result type inference for mhlo.cholesky. 
PiperOrigin-RevId: 448108980
 2022-05-11 16:18:03 -07:00
jax authors
2340725ecc Merge pull request #10669 from mattjj:fix-negative-argnums 
PiperOrigin-RevId: 448105356
 2022-05-11 16:01:02 -07:00
Matthew Johnson
0b841cf35b make core_test.py pass with core.call 2022-05-11 15:45:40 -07:00
Lukas Geiger
f13b69c41d Avoid generating trivial gathers when reversing array 2022-05-11 23:16:40 +01:00
jax authors
5bce8088e4 Merge pull request #10671 from LionOfJewdah:retry_loop_fix 
PiperOrigin-RevId: 448090523
 2022-05-11 14:52:43 -07:00
David Silverstone
4a631886c6 Fixed a typo in cloud_tpu_init.py 2022-05-11 21:19:28 +00:00
Matthew Johnson
7e241b682d improve partial_eval_jaxpr_custom 
* add caching via weakref_lru_cache
* add inst_in argument (needed for fixedpoints for loop primitives, in
  follow-up PR), update callers not to over-instantiate inputs (previously I
  had used a convention where call primitives would just stage out eqns with
  all inputs instantiated, for expediene)
* add ensure_out_unknowns and ensure_out_inst arguments, analogues of
  `instantiate` on e.g. partial_eval_jaxpr, jvp_jaxpr, etc (also neede for
 fixpoints of loop primitives)
* better dce in remat_partial_eval (e.g. prune unused residuals)
 2022-05-11 13:20:23 -07:00
Peter Hawkins
705e241409 Change non-array arguments to jax.lax.linalg functions to be keyword-only arguments. 
PiperOrigin-RevId: 448066207
 2022-05-11 13:06:54 -07:00
jax authors
d092d6305f Merge pull request #10668 from sharadmv:custom-interpreter-update 
PiperOrigin-RevId: 448058058
 2022-05-11 12:31:13 -07:00
jax authors
43467f9873 Merge pull request #10657 from sharadmv:effects-design-note 
PiperOrigin-RevId: 448051701
 2022-05-11 12:03:05 -07:00
Peter Hawkins
590b9161fe Add a sort_eigenvalues option to lax.linalg.eigh(). 
An upcoming change to add a more scalable QDWH-based TPU symmetric eigendecomposition requires that we can obtain the TPU eigenvalues unsorted. The option already exists in XLA, so we simply need to plumb it through to the lax primitive.

PiperOrigin-RevId: 448047584
 2022-05-11 11:46:03 -07:00
Sharad Vikram
2bce098f15 Add sequencing effects design note 2022-05-11 11:19:58 -07:00
Matthew Johnson
28672970bb fix grad(..., argnums=-1), regressed in #10453 2022-05-11 11:19:22 -07:00
Sharad Vikram
aca9dc6c6f Update custom interpreter tutorial 2022-05-11 11:11:47 -07:00
Tianjian Lu
52ad3e6682 [linalg] Extend svd test coverage to input of zero and near-zero elements. 
PiperOrigin-RevId: 448019846
 2022-05-11 09:59:51 -07:00
jax authors
1b0be5095a Merge pull request #10658 from ROCmSoftwarePlatform:rocm_unit_test_enablement 
PiperOrigin-RevId: 448010649
 2022-05-11 09:18:47 -07:00
Filippo Vicentini
a076d349d9
Make CUDA install instructions copy-pasteablee 2022-05-11 17:57:26 +02:00
Yash Katariya
46d034baab Add the nightly dev version to __version__ of jaxlib. 
PiperOrigin-RevId: 448001375
 2022-05-11 08:35:16 -07:00
Eugene Burmako
90f926ac6b [MHLO] Migrate mhlo.all_reduce to HLO_CompatibleOperandsAndResultType 
This runs into the currently unsupported feature in Python bindings which prevents it from taking advantage of the type inference functionality provided by HLO_CompatibleOperandsAndResultType.

PiperOrigin-RevId: 447844880
 2022-05-10 15:39:18 -07:00
Yash Katariya
dfb2caf31e Add nightly __version__ string if building jaxlib nightly 
PiperOrigin-RevId: 447822974
 2022-05-10 14:05:35 -07:00
Sharad Vikram
7d27343506 Attach keepalive to pmap executable 
PiperOrigin-RevId: 447810195
 2022-05-10 13:14:24 -07:00
Tianjian Lu
48f47c36c4 [linalg] Fix a bug in computing derivatives of scipy.special.lpmn. 
PiperOrigin-RevId: 447807140
 2022-05-10 13:02:58 -07:00
Rohit Santhanam
8d9f17df19 Disabled one and enabled several unit tests for ROCm. 2022-05-10 19:47:26 +00:00
Roy Frostig
a62ca21b15 use upgrade option in defining our latest upgrade flag 
PiperOrigin-RevId: 447783540
 2022-05-10 11:25:23 -07:00
jax authors
b826328225 Merge pull request #10649 from onnoeberhard:patch-1 
PiperOrigin-RevId: 447777762
 2022-05-10 11:05:19 -07:00