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

16376 Commits

Author SHA1 Message Date
Roy Frostig
b318f56928 generate name, module, and doctring for functions output from jit. 2018-12-16 13:04:29 -08:00
Matthew Johnson
ea08ecd5f0 add promote_dtypes logic to tensordot 2018-12-16 11:33:57 -08:00
Matthew Johnson
2e20a60916 add tensordot 2018-12-15 21:59:18 -08:00
Matthew Johnson
bfe653c6b0 Tracer.__len__ should reflect on abstract value 
This old implementation, which was meant to be revised but which we
forgot about, caused a surprising slowdown: if x were a traced array of
size 50000, evaluating len(x) would create 50000 traced temporary
objects, which led to a lot of overhead! That came up in our
implementation of jax.random.shuffle, which happened to call len()
instead of x.shape[axis] (even though it should have been using x.size
anyway, according to tjablin@'s code that it's based on).
 2018-12-15 20:07:10 -08:00
Peter Hawkins
06b7e54c02 Fix bug in "economic" mode in jax.scipy.linalg.qr where it returned the full decomposition. 2018-12-15 10:52:10 -05:00
Peter Hawkins
13a135d424 Implement lower=False case for scipy.linalg.cholesky. 
Remove np.linalg.{dot,matmul,trace}, because these aren't part of the numpy API. I had previously misinterpreted the np.linalg documentation to mean that they also existed in that module.
 2018-12-15 10:22:42 -05:00
Matthew Johnson
13b8e21a1c squash conv grad bug introduced in 0d64aea 
(loudly errored, didn't produce silently incorrect results!)
 2018-12-14 18:40:50 -08:00
Matthew Johnson
6de5c8a698 add test-running instructions (fixes #67) 2018-12-14 16:48:08 -08:00
Matthew Johnson
c268929f2d add 'dtype' arg to np.std, add test coverage 2018-12-14 16:22:51 -08:00
sschoenholz
5d6ebba2a0
Fixed argument order in call to var from std. 2018-12-14 11:58:03 -08:00
Matthew Johnson
b164d318fb reduce_and / reduce_or monoid reducer primitives 
The parent commit reused reduce_min / reduce_max on booleans, which is
formally equivalent but preserves less information when lowering to XLA.
 2018-12-14 08:42:02 -08:00
Matthew Johnson
693365c239 np.all and np.any should lead to monoid reducers 
fixes #108
 2018-12-14 08:07:12 -08:00
Peter Hawkins
14acd1dbfa
Merge pull request #110 from hawkinsp/master 
Implement np.linalg.inv using a QR decomposition.
 2018-12-13 21:12:22 -05:00
Peter Hawkins
23525bd9b5 Add scipy.linalg.inv as well. Simplify the QR call in np.linalg.inv. 2018-12-13 21:02:24 -05:00
Matthew Johnson
9b645364c9 tests depend on scipy, tweaks 2018-12-13 16:40:03 -08:00
Peter Hawkins
3aad9b68f6 Implement np.linalg.inv using a QR decomposition. 
An LU decomposition would probably be preferable; we can switch the implementation when we have an LU decomposition.

Fixes #44.
 2018-12-13 19:28:05 -05:00
Peter Hawkins
bc9e157459
Merge pull request #107 from hawkinsp/master 
Make JAX flake8-clean.
 2018-12-13 16:17:20 -05:00
Peter Hawkins
0d4eb6c1e1 Make JAX flake8-clean. 
Fixes #1.
 2018-12-13 15:29:39 -05:00
Matthew Johnson
c5d6c9f09f Merge remote-tracking branch 'origin/master' 2018-12-13 11:56:23 -08:00
Matthew Johnson
bf5a30cdbd fix warning 2018-12-13 11:55:11 -08:00
Matthew Johnson
a285017110 fix failing tests (misc small bugs) 2018-12-13 11:52:41 -08:00
Peter Hawkins
3561b432c2 Add Cholesky, QR, and Triangular solve implementations. 
* Adds lax.{cholesky,triangular_solve,qr}. Adds a JVP for Cholesky.
* Adds a transpose rule for add_p, needed by the Cholesky JVP.
* Adds np.linalg.{cholesky,qr,dot,matmul,trace}.
* Adds scipy.linalg.{cholesky,qr,solve_triangular,tril,triu}.

Pair programmed with mattjj.
 2018-12-13 13:03:08 -05:00
Matthew Johnson
54bceee9e1 make num_generated_cases also settable by env var 2018-12-13 08:56:40 -08:00
Matthew Johnson
579665b1b0 add set_printoptions (same as onp version) 2018-12-13 07:28:59 -08:00
Peter Hawkins
22f9d50be0 Implement np.trace. 
Also add code to forward np.{iinfo,finfo,issubdtype,integer} to their original numpy equivalents.

Fixes #29.
 2018-12-13 08:44:27 -05:00
Matthew Johnson
3fba60fcd5
Merge pull request #100 from vishwakftw/log2-log10-exp2-logaddexp-logaddexp2 
Add log2, log10, logaddexp2, exp2 to lax_numpy
 2018-12-13 05:09:39 -08:00
vishwakftw
52d329c621 Add log2, log10, logaddexp2, exp2 to lax_numpy 
I also took this opportunity to add some missing decorators
 2018-12-13 12:57:49 +05:30
Matthew Johnson
1f42d980b8 rename ResNet50Test -> ExamplesTest, remove some dots 2018-12-12 20:28:54 -08:00
Matthew Johnson
77d6fb4c01 transpose shouldn't transpose with identity perm 2018-12-12 20:28:54 -08:00
Matthew Johnson
538850e271 add misc numpy ops (c.f. #70) 2018-12-12 20:28:54 -08:00
Peter Hawkins
9c63c0ff0a Fix dtype semantics of numpy reductions to more closely match that of numpy. 
Add tests for the dtype argument to reductions.
 2018-12-12 21:53:18 -05:00
Peter Hawkins
7530f56809 Implement np.{diag,diagonal}. 2018-12-12 17:54:27 -05:00
Matthew Johnson
986514eeda
Merge pull request #71 from vishwakftw/scipy-stats-more-dists 
Add a few popular distributions to jax.scipy.stats
 2018-12-12 10:16:06 -08:00
Matthew Johnson
af8223a2ad
Merge pull request #90 from hawkinsp/master 
Add implementations of np.{tri,tril,triu}.
 2018-12-12 09:17:40 -08:00
Peter Hawkins
3c11d4f780 Add implementations of np.{tri,tril,triu}. 2018-12-12 12:05:49 -05:00
Matthew Johnson
d3ec0b23c5 fix miscellaneous bugs, incl. complex abs grad 2018-12-12 12:03:57 -05:00
vishwakftw
0e6a44f1d2 Add tests, fix small bugs in the logpdf functions 2018-12-12 20:41:48 +05:30
Dougal Maclaurin
9b835a180f Added some docstrings to top-level transformations 2018-12-12 09:47:49 -05:00
Peter Hawkins
dad72f94da Add missing license header to Python files. 2018-12-12 08:45:06 -05:00
Matthew Johnson
89349e5e6d fix transpose issue in jacfwd and jacrev 2018-12-11 16:24:20 -08:00
Matthew Johnson
e788539e0a add concatenate batching rule (c.f. #54) 2018-12-11 16:11:54 -08:00
Matthew Johnson
97589a3d03 add batching rule for lax.pad (c.f. #54) 2018-12-11 16:11:54 -08:00
Matthew Johnson
cbd9ace425 change vmap api to be curried (closes #78) 2018-12-11 12:52:09 -08:00
Matthew Johnson
9cd6027979 add python built-in complex type to array types 
fixes #74
 2018-12-11 12:14:57 -08:00
Roy Frostig
2b9a8bc13b
Merge pull request #77 from google/rot90 
add rot90 and flip, adjust testOp test selection
 2018-12-11 11:44:15 -08:00
Matthew Johnson
5a1aeca96c add rot90 and flip, adjust testOp test selection 
closes #55
 2018-12-11 11:33:15 -08:00
Peter Hawkins
0f9a69fb81
Merge pull request #72 from google/unimplemented-funcs 
Adding unimplemented functions more numpy modules
 2018-12-11 13:58:07 -05:00
Matthew Johnson
7198e09465 enable skip_checks for merging to master 2018-12-11 13:22:07 -05:00
Matthew Johnson
ffecf2ccc5 fix bugs in zeros_like and complex transpose 2018-12-11 13:22:07 -05:00
Matthew Johnson
0d64aea6bb clean up conv dimension_numbers handling 2018-12-11 13:22:07 -05:00