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

699 Commits

Author SHA1 Message Date
Lukas Geiger
b4b02cb7c3 Add jnp.resize 2021-05-21 09:25:49 +01:00
Lukas Geiger
a5a1f62c5c Improve docs for jnp.round and export jnp.round_ 2021-05-21 01:43:56 +01:00
Paul Nguyen
bbcaec4a3a Initial implementation of jax.numpy.poly 
This is an initial jax.numpy.poly implementation. It is tested by testPoly in the tests/lax_numpy_test.py test file.
 2021-05-20 13:51:14 -05:00
Nicholas Junge
e985e86808 Implement jnp.r_ and jnp.c_ 2021-05-01 01:05:22 +02:00
Lukas Geiger
a613ce12a3 Export jnp.broadcast_shapes as user facing function 2021-04-30 19:32:51 +02:00
Peter Hawkins
4409e03c3b Remove workaround for NumPy < 1.14. 
Following NEP 29, we don't need to support NumPy older than 1.17 these days (although in practice we also support 1.16).

Remove accidental export jnp.numpy_version.
 2021-04-20 15:27:44 -04:00
minoring
4c67dd1f48 Implement jnp.ogrid 
Related to #5850
 2021-04-06 08:37:07 +09:00
minoring
a2209aa509 Implement jnp.mgrid 
Related to #5850
 2021-04-02 12:38:20 +09:00
Jake VanderPlas
737e4796cd Initial implementation of jnp.delete 2021-03-16 17:05:23 -07:00
Jake VanderPlas
fcaf8bf4f8 Add test that docstring wrapping didn't fail 2021-03-04 14:16:43 -08:00
jax authors
932e118dd2 Merge pull request #5665 from ashutoshvarma:numpy-setxor1d 
PiperOrigin-RevId: 359347458
 2021-02-24 12:43:14 -08:00
Ashutosh Varma
c223ef9d07 add support for setxor1d 2021-02-24 23:28:42 +05:30
Jonathan Terhorst
4c202ad222 implement np.polyint (#70) 2021-02-18 11:08:41 -05:00
Ashutosh Varma
108d078c5f add support for union1d 
add union1d in jax.numpy which closely follows
numpy implementaion
 2021-02-06 22:19:16 +05:30
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