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
11,145 Commits 506 Branches 263 Tags
Commit Graph

11145 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
jax authors
ef2efec649 Merge pull request #10069 from jakevdp:devicearray-copy 
PiperOrigin-RevId: 438292130
 2022-03-30 07:01:19 -07:00
jax authors
17fc5bd02e Merge pull request #9290 from fehiepsi:named 
PiperOrigin-RevId: 438290209
 2022-03-30 06:54:10 -07:00
jax authors
4ac8eb055b Merge pull request #9645 from yalaudah:main 
PiperOrigin-RevId: 438290146
 2022-03-30 06:49:16 -07:00
jax authors
02c2eb59f9 Merge pull request #9682 from jblespiau:changelist/430664138 
PiperOrigin-RevId: 438285304
 2022-03-30 06:22:20 -07:00
jax authors
f8cddf0eca Merge pull request #9689 from fehiepsi:image 
PiperOrigin-RevId: 438282642
 2022-03-30 06:05:30 -07:00
jax authors
9da5f4e793 Merge pull request #10074 from jakevdp:jit-partial-doc 
PiperOrigin-RevId: 438193490
 2022-03-29 20:02:17 -07:00
Aden Grue
8884ce5b98 Migrate 'jaxlib' CPU custom-calls to the status-returning API 
PiperOrigin-RevId: 438165260
 2022-03-29 17:14:14 -07:00
Jake VanderPlas
f4b64f48f4 doc: add examples of using partial with jit 2022-03-29 15:43:58 -07:00
jax authors
b31cf89e48 Merge pull request #10072 from jakevdp:fromiter 
PiperOrigin-RevId: 438141629
 2022-03-29 15:28:36 -07:00
jax authors
23c783af88 Merge pull request #10068 from hawkinsp:docs 
PiperOrigin-RevId: 438131830
 2022-03-29 14:49:47 -07:00
Jake VanderPlas
4f6ea7b296 docs: use intersphinx links for wrapped functions 2022-03-29 14:43:59 -07:00
Jake VanderPlas
fbfc3d8edf Better error messages for jnp.fromiter and jnp.fromfile 2022-03-29 14:30:32 -07:00
Yash Katariya
e08bc27bf0 Speed up GDA initialization by making local_shards lazy and hiding checks behind config.jax_enable_checks flag. 
PiperOrigin-RevId: 438115859
 2022-03-29 13:51:11 -07:00
jax authors
085d3901fd Merge pull request #9369 from froystig:custom-vmap-outputs 
PiperOrigin-RevId: 438111265
 2022-03-29 13:35:49 -07:00
Roy Frostig
b2de101be7 require consistent output structure in custom vmap rules 
... not always a sequence.
 2022-03-29 12:28:04 -07:00
jax authors
4b4010dbd8 Merge pull request #10049 from jakevdp:from-funcs 
PiperOrigin-RevId: 438086962
 2022-03-29 12:03:44 -07:00
jax authors
90629596a4 Merge pull request #10064 from ayaka14732:patch-2 
PiperOrigin-RevId: 438069965
 2022-03-29 11:00:24 -07:00
Jake VanderPlas
093b7032a8 Implement jnp.from* array creation functions 2022-03-29 10:52:47 -07:00
Jake VanderPlas
de9a948d1c make device_array.copy() return a device array 2022-03-29 10:33:29 -07:00
Peter Hawkins
150235ca5b Add note about enabling developer mode on Windows. 
Issue 4843.
 2022-03-29 12:32:42 -04:00
Ayaka Mikazuki
2799bb3cde
[doc] Fix typo 2022-03-29 21:29:51 +08:00
Yash Katariya
a68b0f3a0a Fix the mesh_axes of benchmarks to be Pspecs 
PiperOrigin-RevId: 437932954
 2022-03-28 21:51:01 -07:00
jax authors
d69c7b3c21 Merge pull request #10054 from froystig:upgrade-flags 
PiperOrigin-RevId: 437931412
 2022-03-28 21:39:15 -07:00
Tianjian Lu
e5d7f65b6a [sparse] Change call signature of bcoo primitive wrappers. 
PiperOrigin-RevId: 437923482
 2022-03-28 20:42:52 -07:00
Roy Frostig
e0ddd43933 add custom_vjp upgrade flag to jax module and fix doc rendering for upgrade flags 2022-03-28 20:36:35 -07:00
jax authors
ab8dd4e623 Merge pull request #10048 from froystig:custom-vjp-by-jvp 
PiperOrigin-RevId: 437891536
 2022-03-28 17:13:44 -07:00
Roy Frostig
941b221de6 refine occurrences of the "forward-mode disallowed" error under new custom_vjp 
To achieve this, this change introduces a general primitive
(`unreachable`) that allows staging but not evaluation. We use it to
allow for the `custom_transpose`d tangent function to be stageable,
but to ensure it errs if it is evaluated (which only happens due to
forward-mode).

Co-authored-by: Matthew Johnson <mattjj@google.com>
 2022-03-28 16:50:47 -07:00
jax authors
4761b1ea0b Merge pull request #10053 from mattjj:remove-unused-arg 
PiperOrigin-RevId: 437880655
 2022-03-28 16:20:45 -07:00
Matthew Johnson
65d05984f0 remove unused arg in call param jvp updater 2022-03-28 15:40:49 -07:00
Roy Frostig
52a1618c3a introduce partly-working new custom_vjp implementation behind a flag 
... by reduction to `custom_jvp` and `custom_transpose`
 2022-03-28 14:13:29 -07:00
Roy Frostig
1d9a6c1f82 introduce an upgrade flag for custom_vjp overhaul (still experimental) 
Along the way, make "boolean upgrade flag" a common option for current
and future re-use.
 2022-03-28 14:10:08 -07:00
jax authors
5dc068f54b Merge pull request #10042 from yotarok:fix_fft_helper 
PiperOrigin-RevId: 437804929
 2022-03-28 11:12:15 -07:00
jax authors
b0f3cdcdda Merge pull request #10043 from gnecula:poly_reshape 
PiperOrigin-RevId: 437804782
 2022-03-28 11:07:09 -07:00
jax authors
1cd1b11174 Merge pull request #10033 from LenaMartens:changelist/433197432 
PiperOrigin-RevId: 437728213
 2022-03-28 05:13:34 -07:00
George Necula
4002f6b6dd [shape_poly] Change internal representation of parametric shapes 
The parametric shapes are implemented as polynomials in terms of
dimension variables. In the previous implementation the representation
is as a `dict`. This happens to confuse a test in `jnp.reshape` because
the polynomial appears to be an iterable.

I removed the inheritance from `dict`, so that a dimension polynomial
is harder to confuse with other types.

Fixes: #9975
 2022-03-28 14:06:36 +03:00
Lena Martens
d72687c990 Checkify: add way to embed runtime info in error messages. 2022-03-28 11:55:34 +01:00
Yotaro Kubo
2e70177385 Fix a bug in fft helper appears when nperseg=1. 2022-03-28 14:51:54 +09:00
jax authors
c6debb5921 Merge pull request #10041 from froystig:custom-transpose-recursion2 
PiperOrigin-RevId: 437665703
 2022-03-27 21:41:03 -07:00
Roy Frostig
a6a43e2715 allow for recursive uses of custom_transpose 
Co-authored-by: Matthew Johnson <mattjj@google.com>
 2022-03-26 12:09:15 -07:00
jax authors
f7df3ee9c4 Merge pull request #10037 from mattjj:improve-remat-caching2 
PiperOrigin-RevId: 437397287
 2022-03-25 21:22:55 -07:00
Matthew Johnson
78cf4df21b improve remat transpose caching (cf. #9661) 2022-03-25 16:33:46 -07:00
jax authors
563e0c6ba8 Merge pull request #10034 from mattjj:improve-remat-caching 
PiperOrigin-RevId: 437356036
 2022-03-25 16:12:11 -07:00
jax authors
330aded7e0 Merge pull request #10029 from jakevdp:rm-lax-numpy 
PiperOrigin-RevId: 437347088
 2022-03-25 15:24:14 -07:00
Matthew Johnson
bd765fecb5 improve caching of jax.remat 
See #9661 for discussion
 2022-03-25 15:15:30 -07:00
Jake VanderPlas
f4d240c036 Remove lax_numpy from jax.numpy namespace 
This is a private module that was inadvertently exported in the past.
 2022-03-25 15:02:45 -07:00
jax authors
fe035f56fc Merge pull request #10019 from hyeontaek:transfer-guard-doc 
PiperOrigin-RevId: 437319396
 2022-03-25 13:16:12 -07:00
Hyeontaek Lim
87671f4d85 Add a documentation for the JAX transfer guard 2022-03-25 19:50:12 +00:00
jax authors
f57e78e240 Merge pull request #9766 from ayaka14732:patch-1 
PiperOrigin-RevId: 437297267
 2022-03-25 11:39:41 -07:00
jax authors
21095b04dd Merge pull request #9986 from ROCmSoftwarePlatform:fix_hipsparse_spmm_coo 
PiperOrigin-RevId: 437292738
 2022-03-25 11:18:29 -07:00
Reza Rahimi
8cd02946b5 Fix for hipsparse in ROCm. 2022-03-25 17:41:42 +00:00