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

685 Commits

Author SHA1 Message Date
Peter Hawkins
3ac809ede3 [JAX] Move jax.util to jax._src_util. 
PiperOrigin-RevId: 351234602
 2021-01-11 14:21:07 -08:00
Alex Dragan
b7fee456e6 [impl] Add support for setdiff1d 2020-10-21 20:17:12 +03:00
Peter Hawkins
3ddd3905a4 Move jax.third_party to jax._src.third_party. 
PiperOrigin-RevId: 337675377
 2020-10-17 11:43:07 -07:00
Peter Hawkins
aa107cf1f4 Move jax.numpy internals into jax._src.numpy. 2020-10-16 20:35:19 -04:00
Matthew Johnson
e51163af32 only pass hashable values as static args 2020-10-16 13:11:56 -07:00
jax authors
b2db682a08 Merge pull request #4586 from shoyer:precision-tuple 
PiperOrigin-RevId: 337217160
 2020-10-14 18:33:47 -07:00
Stephan Hoyer
706bab7cf4 Support specifying precision as a tuple of precision enums 
This is useful for cases where one of the arguments is known to have
low precision, e.g., on https://github.com/google/jax/pull/4459 where it
would potentially allow for saving half the FLOPs on TPUs when using
convolutions to extract image patches.
 2020-10-14 18:07:37 -07:00
jax authors
76fe1803c2 Merge pull request #4556 from jakevdp:index-fix 
PiperOrigin-RevId: 336900642
 2020-10-13 10:09:06 -07:00
jax authors
f077139178 Merge pull request #4451 from alexminnaar:master 
PiperOrigin-RevId: 336888044
 2020-10-13 09:08:28 -07:00
Matthew Johnson
dbebc9ad04 remove double warning with asarray dtype='int' 2020-10-12 15:00:35 -07:00
Jake VanderPlas
3ca9bba03a BUG: fix indexing error 2020-10-12 13:58:49 -07:00
jax authors
9d8139e0ea Merge pull request #4440 from google:long-line 
PiperOrigin-RevId: 336638751
 2020-10-12 04:49:55 -07:00
Alex Minnaar
0485fd8736 adding histogram2d implementation 2020-10-10 12:03:46 -04:00
Jake VanderPlas
cff6c0b05a Improve errors for failed compilations w/ core.concrete_or_error 2020-10-10 12:03:46 -04:00
Jake VanderPlas
67a77c7ee4 Fix extraneous dtype warning in jnp.mean 2020-10-09 09:13:21 -07:00
Peter Hawkins
a8b67ae409 Switch implementation of jnp.isnan(x) to x != x. 2020-10-07 16:54:13 -04:00
jax authors
efeed6c1a8 Merge pull request #4444 from jakevdp:dtype-check 
PiperOrigin-RevId: 335442628
 2020-10-05 12:09:05 -07:00
jax authors
0e92124c5b Merge pull request #4443 from jakevdp:concrete 
PiperOrigin-RevId: 335180298
 2020-10-03 03:23:55 -07:00
Jake VanderPlas
0328f858e1 Call check_user_dtype on all user dtypes 2020-10-02 17:21:44 -07:00
Jake VanderPlas
492889f4a4 Improve errors for failed compilations w/ core.concrete_or_error 2020-10-02 15:08:21 -07:00
Jake VanderPlas
2fccbbf65b Implement jnp.choose 2020-10-02 13:13:21 -07:00
Roy Frostig
081ba14d89 wrap long line 2020-10-02 11:03:36 -07:00
jax authors
b609040ce6 Merge pull request #4413 from alexminnaar:master 
PiperOrigin-RevId: 334946450
 2020-10-01 19:26:55 -07:00
Alex Minnaar
0893b08a6f histogramdd implemenation 2020-10-01 18:02:33 -04:00
jax authors
69fda9ecb9 Merge pull request #4039 from LenaMartens:changelist/325216264 
PiperOrigin-RevId: 334728148
 2020-09-30 19:25:00 -07:00
jax authors
768b0c1eb7 Merge pull request #4415 from jakevdp:indices-concrete 
PiperOrigin-RevId: 334683008
 2020-09-30 14:44:47 -07:00
Jake VanderPlas
0b23a90496 Call concrete_or_error on indices() argument 2020-09-30 11:44:49 -07:00
jax authors
47da5121ac Merge pull request #4384 from jakevdp:cov-weights 
PiperOrigin-RevId: 334626974
 2020-09-30 10:23:34 -07:00
Akihiro Nitta
06170da69a
Use raise from 2020-09-30 01:20:00 +09:00
Jake VanderPlas
a51a4d91b3 Add test coverage for jnp.cov aweights & fweights 2020-09-28 15:34:57 -07:00
Lena Martens
ecad419cf3 Support grad with integer arguments. 
- Add float0 and set-up at_least_vspace to return float0
values for int/bool primals
- Use Zero to wrap float0 tangents so they're correctly ignored in jvp
rules
- Add float0 handlers to XLA to support jit
- Fix convert_element_type and tie_in jvp rules
 2020-09-28 19:07:04 +01:00
Jake VanderPlas
185590fcbf Use core.concrete_or_error() to improve errors in reductions 2020-09-25 14:18:46 -07:00
jax authors
8d24d0af26 Merge pull request #4368 from jakevdp:unop-check-array 
PiperOrigin-RevId: 333765605
 2020-09-25 10:56:00 -07:00
Jake VanderPlas
ab273acb21 Call _check_arraylike in jax.numpy to improve error messages 
This PR only adds the call to places where non-array inputs currently lead to errors. There
remain a number of other functions where adding this check would lead to potentially breaking
changes; these are deliberately left out of this PR.
 2020-09-22 10:01:55 -07:00
jax authors
d43d5d9035 Merge pull request #4371 from jakevdp:moveaxis-fix 
PiperOrigin-RevId: 332974804
 2020-09-21 17:44:35 -07:00
Jake VanderPlas
2cf8d49f5b jnp.moveaxis: fix bug when axes are integer dtype 2020-09-21 16:32:11 -07:00
jax authors
2cb795e0d9 Merge pull request #4366 from jakevdp:diff-empty 
PiperOrigin-RevId: 332960638
 2020-09-21 16:28:30 -07:00
Jake VanderPlas
80d5f9ce6c jax.numpy: improved errors for invalid inputs to unary ops 2020-09-21 15:26:47 -07:00
Jake VanderPlas
18054e05a8 call _check_arraylike in jnp.diff 2020-09-21 15:00:14 -07:00
Jake VanderPlas
ce1ce9cb27 Implement jnp.array_equiv 2020-09-21 10:37:55 -07:00
Jake Vanderplas
e0af77fbca
Implement jnp.ravel_multi_index() (#4313) 2020-09-17 12:42:22 -07:00
Srijan Saurav
40e20242db
Fix code quality issues (#4302) 
Changes:
- Fix unnecessary generator
- Iterate dictionary directly instead of calling .keys()
- Remove global statement at the module level
- Use list() instead of a list comprehension
- Use with statement to open the file
- Merge isinstance calls
 2020-09-17 09:21:18 -07:00
Julius Kunze
c6b7269480
Support non-fragmenting mask of reshape (#4264) 2020-09-16 23:58:32 -07:00
Jake Vanderplas
e18a973198
implement jnp.apply_over_axes (#4225) 2020-09-16 13:30:08 -07:00
Jake Vanderplas
19201f4b4a
Mention in docstring when function is not implemented (#4297) 2020-09-15 18:16:36 -07:00
Jake Vanderplas
9c393812c4
jnp.divide: remove obsolete condition for py2 behavior (#4286) 2020-09-15 10:58:47 -07:00
Matthew Johnson
2678a4647a
omnistaging on by default (#4038) 2020-09-15 08:06:46 -07:00
Benjamin Chetioui
58a117fe0d
Modifies eig_p and related operations to take advantage of the new jaxlib geev API (#4266) 
* Add options to compute L/R eigenvectors in geev.

The new arguments are by default set to True to ensure backwards
compatibility between jaxlib and jax.

Reformulate eig-related operations based on the new geev API.

* Addressed hawkinsp's comments from google/jax#3882.

Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
 2020-09-15 11:45:15 +03:00
Roman Novak
38b43ef95d
Avoid rank promotion in np.outer (#4276) 2020-09-13 19:20:31 -07:00
Jake Vanderplas
83b4f3b97c
Cleanup: use _canonicalize_axis() utility where possible (#4270) 2020-09-11 16:49:18 -07:00