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
10,633 Commits 504 Branches 263 Tags
Commit Graph

10633 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
jax authors
fb44d7c12f [JAX] Add release note for the graduration of the experimental.ann module. 
PiperOrigin-RevId: 431951602
 2022-03-02 08:58:53 -08:00
jax authors
3766dd2120 Rollback of: 
d09d7b8d1363eab1c14051eb2376e605366537f9 by Jake VanderPlas <jakevdp@google.com>:

Factor-out pieces of lax_numpy.py

PiperOrigin-RevId: 431833044
 2022-03-01 19:39:31 -08:00
Yash Katariya
4755dc3fee Don't compute global shards if local_mesh == global_mesh and just return local_shards. 
PiperOrigin-RevId: 431815962
 2022-03-01 17:48:29 -08:00
Yash Katariya
72cc567c05 Use the new mesh property instead of the private _global_mesh attribute. 
PiperOrigin-RevId: 431815802
 2022-03-01 17:43:40 -08:00
jax authors
d9f82f7b9b [JAX] Move experimental.ann.approx_*_k into lax. 
Updated docs, tests and the example code snippets.

PiperOrigin-RevId: 431781401
 2022-03-01 14:46:33 -08:00
jax authors
1a1bf122d9 [JAX] Fix github lint_and_typecheck for Primitives 
PiperOrigin-RevId: 431776305
 2022-03-01 14:24:59 -08:00
jax authors
c50b9d470a Merge pull request #9736 from mattjj:test-issue-9731 
PiperOrigin-RevId: 431765284
 2022-03-01 13:39:09 -08:00
Matthew Johnson
4075b81e00 add regression test for #9731 2022-03-01 13:10:28 -08:00
jax authors
d6812b24dc Merge pull request #9732 from jakevdp:jnp-copy 
PiperOrigin-RevId: 431754341
 2022-03-01 12:50:17 -08:00
Jake VanderPlas
ed2550999f implement jnp.copy 2022-03-01 11:56:36 -08:00
Yash Katariya
11664f8aee Make the global_mesh that GDA has public 
PiperOrigin-RevId: 431739364
 2022-03-01 11:45:54 -08:00
jax authors
5b309880bd Merge pull request #9724 from jakevdp:refactor-lax-numpy 
PiperOrigin-RevId: 431736367
 2022-03-01 11:33:39 -08:00
Jake VanderPlas
d09d7b8d13 Factor-out pieces of lax_numpy.py 2022-03-01 09:40:37 -08:00
Peter Hawkins
cffe9978fb Handle jaxpr constants correctly in MLIR lowering of conditional branches. 
Add some dynamic type checks and type annotations to catch this kind of problem sooner.

There's no test case, because I'm not entirely sure how to make a test case for this. In fact, I'm not even sure it's legal for a conditional branch to have non-empty constants. We'll dig into that separately.

PiperOrigin-RevId: 431697808
 2022-03-01 08:56:08 -08:00
jax authors
c7508d1f2d Merge pull request #9721 from jakevdp:poisson-nan 
PiperOrigin-RevId: 431505317
 2022-02-28 12:59:08 -08:00
Yash Katariya
d0cc3395e8 Add block_until_ready method to GDA 
PiperOrigin-RevId: 431504594
 2022-02-28 12:54:23 -08:00
Jake VanderPlas
2c2773a5f1 jax.random.poisson: fix corner cases 2022-02-28 12:10:47 -08:00
jax authors
92cb865b3c Merge pull request #9718 from google:dependabot/github_actions/actions/setup-python-3 
PiperOrigin-RevId: 431490872
 2022-02-28 11:54:29 -08:00
Peter Hawkins
c339330bc1 [XLA:CPU] Relax test tolerances for tests using XLA:CPU. 
An upcoming change to XLA:CPU will disable reassociation on floating point operators by default which is an unsound fast math optimization. This change is being made to fix numerical errors in softmax computations caused by reassocation. After that change, we will enable reassociation only in reduction operators where it is very important for performance and the XLA operator contract allows that.

Since this change alters the order of operations, it may cause small numerical changes leading to test failures. This change relaxes test tolerances to make tests pass.

PiperOrigin-RevId: 431453240
 2022-02-28 09:26:54 -08:00
dependabot[bot]
680c06ddc4
Bump actions/setup-python from 2 to 3 
Bumps [actions/setup-python](https://github.com/actions/setup-python) from 2 to 3.
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/v2...v3)

---
updated-dependencies:
- dependency-name: actions/setup-python
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
 2022-02-28 17:11:41 +00:00
jax authors
043561ae02 Merge pull request #9711 from jblespiau:changelist/431385993 
PiperOrigin-RevId: 431447036
 2022-02-28 08:57:42 -08:00
jax authors
0b9ae982cb Merge pull request #9701 from jakevdp:fix-sharp-bits 
PiperOrigin-RevId: 431435194
 2022-02-28 07:56:57 -08:00
Jean-Baptiste Lespiau
25472c238b Remove an unnecessary condition. 
This makes sure jit, xla_computation and pjit share the same logic for processing static arguments.
 2022-02-28 11:34:01 +01:00
jax authors
c66fbeb99b Merge pull request #9193 from froystig:xla-executable-attr 
PiperOrigin-RevId: 431102224
 2022-02-25 23:29:04 -08:00
Roy Frostig
d636e74626 make xla_executable a property, consistent across executable types 
Also test IR and executable-related methods of `Lowered` and
`Compiled`.
 2022-02-25 19:05:44 -08:00
jax authors
d583f87618 Merge pull request #9706 from samuela:samuela/mkl-fix 
PiperOrigin-RevId: 431028480
 2022-02-25 14:19:17 -08:00
Samuel Ainsworth
bf59b7d872 Relax tolerances slightly for MKL. 
Fix https://github.com/google/jax/issues/9705.
 2022-02-25 22:02:55 +00:00
jax authors
7ec04c8311 Merge pull request #9702 from jakevdp:fix-load 
PiperOrigin-RevId: 431016627
 2022-02-25 13:23:11 -08:00
Jake VanderPlas
1b01865b89 BUG: return numpy arrays for jnp.load() with unsupported dtypes 2022-02-25 09:27:42 -08:00
Jake VanderPlas
74e7fdfd4c DOC: update sharp bits info about fori_loop differentiability 2022-02-25 09:15:55 -08:00
jax authors
d5a1c64d13 Merge pull request #9623 from rsuderman:FixIreeJax 
PiperOrigin-RevId: 430834942
 2022-02-24 17:51:09 -08:00
jax authors
6914bae32b Merge pull request #9641 from rsuderman:FixDefaultDevice 
PiperOrigin-RevId: 430828922
 2022-02-24 17:13:48 -08:00
jax authors
ffcf477345 Merge pull request #9684 from jakevdp:fix-changelog 
PiperOrigin-RevId: 430827769
 2022-02-24 17:07:02 -08:00
Jake VanderPlas
51727033b8 Remove duplicate changelog entry 2022-02-24 08:18:30 -08:00
jax authors
d2834b64df Merge pull request #9658 from tornikeo:docs-copybutton 
PiperOrigin-RevId: 430705653
 2022-02-24 08:12:41 -08:00
jax authors
8372b98c48 [JAX] Move ann.ann_recall back to tests. 
The function is simple enough for users to implement their own on the host.

PiperOrigin-RevId: 430696789
 2022-02-24 07:23:17 -08:00
jax authors
cdd0c68877 Merge pull request #9619 from hawkinsp:index2 
PiperOrigin-RevId: 430692755
 2022-02-24 07:02:34 -08:00
jax authors
f2c1743972 [JAX] Update comments and documents for ANN. 
Copied some of the examples from the module to the functions to make the
migration to jax.lax easier, because only the functions will be copied over.

PiperOrigin-RevId: 430692713
 2022-02-24 06:57:56 -08:00
Peter Hawkins
f51a05a889 Remove jax.ops.index... functions. 
These functions have been deprecated and have issued a DeprecationWarning since jax 0.2.22 in October 2021.
 2022-02-24 09:36:28 -05:00
jax authors
3948fde842 Merge pull request #9052 from jpuigcerver:main 
PiperOrigin-RevId: 430680329
 2022-02-24 05:37:02 -08:00
Yash Katariya
a9a827e28c Add .device() method to _DeviceArray 
PiperOrigin-RevId: 430551695
 2022-02-23 15:22:37 -08:00
Yash Katariya
98e114da4f Rename unmapped_local_out_avals to out_avals since it can contain global avals (via GDA) as well as local avals. 
PiperOrigin-RevId: 430539281
 2022-02-23 14:27:00 -08:00
Yash Katariya
e2834d89e1 Fix the gpu tests that were failing with Future warning 
PiperOrigin-RevId: 430532523
 2022-02-23 13:58:28 -08:00
jax authors
c041694538 Merge pull request #8395 from sharadmv:name-stack-mechanism 
PiperOrigin-RevId: 430506453
 2022-02-23 12:00:20 -08:00
Yash Katariya
687a7630ee Deprecate maps.mesh and replace it with maps.Mesh. 
PiperOrigin-RevId: 430489855
 2022-02-23 10:47:06 -08:00
Sharad Vikram
1b79caa6bd Add separate mechanism for threading name stacks to the lowering 2022-02-23 09:59:09 -08:00
Yash Katariya
e96b91d405 Make the name pjit appear in xprof and mhlo module name. Before everything was xmap even for pjit. 
PiperOrigin-RevId: 430477182
 2022-02-23 09:55:12 -08:00
jax authors
7edd964cc3 Merge pull request #9663 from mattjj:checkify-scan-debug 
PiperOrigin-RevId: 430462126
 2022-02-23 08:43:09 -08:00
lenamartens
45d3ddda31 Fix tests and handle cond consts. 2022-02-23 16:11:09 +00:00
Matthew Johnson
4b1d0a466b fixing scan and other control flow 
Co-authored-by: Lena Martens <lenamartens@google.com>
 2022-02-23 15:39:22 +00:00