Peter Hawkins
95135377d0
Add implementations of np.{meshgrid,linspace,logspace,geomspace,diag_indices} that forward to the usual numpy implementation.
2018-12-20 08:28:36 -05:00
Matthew Johnson
449da4cdb5
Merge pull request #147 from google/einsum
...
fix einsum bugs, add test cases
2018-12-19 17:12:16 -08:00
Matthew Johnson
6bb9609fb8
disable test, py3 opt_einsum nondeterministic bug?
2018-12-19 16:58:31 -08:00
Matthew Johnson
6a138202ef
fix several einsum bugs
2018-12-19 16:15:43 -08:00
Matthew Johnson
569db1698c
Merge pull request #148 from google/solve-triangular-jvp
...
Solve triangular jvp rule
2018-12-19 15:34:11 -08:00
Dougal Maclaurin
e6b23dd2b7
Fixed triangular_solve_jvp_rule for transpose_a=True case
2018-12-19 17:47:56 -05:00
Matthew Johnson
9c722db373
einsum: update id strings after moving batch dims
2018-12-19 10:59:03 -08:00
Matthew Johnson
9a68bce567
add comment marking a bug
2018-12-19 10:42:40 -08:00
Matthew Johnson
43e77acca5
fix select transpose rule
2018-12-19 09:40:40 -08:00
Matthew Johnson
dc1d0c260a
always fall back to onp.arange for now
2018-12-19 09:21:30 -08:00
Matthew Johnson
6a9952a939
jax.numpy.arange should fall back to onp.arange
...
fixes #145
2018-12-19 09:07:04 -08:00
Matthew Johnson
c56f43f2a5
Merge pull request #144 from google/einsum
...
fix einsum tensor product logic (fixes #37 )
2018-12-19 08:27:16 -08:00
Dougal Maclaurin
87922fdf13
Generalized make_jaxpr to handle python containers
2018-12-19 10:59:13 -05:00
Matthew Johnson
997c9c5a50
fix einsum tensor product logic ( fixes #37 )
...
The error was that `lhs_names` and `rhs_names` included `batch_names` as
prefixes, but the reshaping logic was written as if they did not include
batch_names (and so batch_names had to be prepended).
2018-12-19 07:59:00 -08:00
Peter Hawkins
a154b9c691
Merge pull request #143 from hawkinsp/astype
...
Fix implementation of ndarray.astype method, add a test.
2018-12-19 09:57:26 -05:00
Peter Hawkins
d3bb93f82a
Fix implementation of ndarray.astype method, add a test.
...
Previously we were creating an operator named __astype__, which isn't a thing in numpy.
2018-12-19 09:28:08 -05:00
Peter Hawkins
dbab47fdb0
Implement np.inner and np.outer.
2018-12-19 08:58:07 -05:00
Matthew Johnson
0a9ee106b3
implement triangular solve lhs jvp (w/ @froystig)
2018-12-18 23:39:46 -08:00
Matthew Johnson
d7745dd9af
actually fix py3 str translate
2018-12-18 23:20:10 -08:00
Matthew Johnson
061d033c2b
add jit around einsum
2018-12-18 23:20:10 -08:00
Matthew Johnson
6261ef729a
more einsum improvements (complete?)
2018-12-18 23:20:10 -08:00
Matthew Johnson
d8388e2d80
complete support for two-operand einsum
2018-12-18 23:20:10 -08:00
Matthew Johnson
fdde6841e6
add support for two-operrand cases
2018-12-18 23:20:10 -08:00
Matthew Johnson
13a0e1168e
fix broadcasted eye bug, enable more einsum
2018-12-18 23:20:10 -08:00
Matthew Johnson
6a71e9d6ec
start drafting an einsum implementation
2018-12-18 23:20:09 -08:00
Matthew Johnson
a18e3f27ac
add tests for device constants
2018-12-18 22:45:34 -08:00
Matthew Johnson
20ca0bd733
add cutoff for materialize-and-xfer vs build-on-device
...
see https://github.com/google/jax/pull/140#issuecomment-448433620
2018-12-18 17:19:47 -08:00
Matthew Johnson
78a4240581
add fix from einsum branch
2018-12-18 09:18:14 -08:00
Matthew Johnson
52c6eac3de
use lax.tie_in in jax.random for better consts
2018-12-18 09:16:59 -08:00
Matthew Johnson
bdc9e92f94
remove full_p to improve power:weight
2018-12-18 09:16:59 -08:00
Matthew Johnson
1ae1ae17a2
add EyeConstant, new np.eye and np.array code
2018-12-18 09:16:59 -08:00
Matthew Johnson
dfc25a06d9
add IotaConstant (untested)
2018-12-18 09:16:59 -08:00
Matthew Johnson
f971415218
add tie_in and full primitives (constant creation)
2018-12-18 09:16:59 -08:00
Matthew Johnson
25cf9358d1
Merge pull request #131 from google/mean-kwargs
...
fix mean/var/std kwargs (closes #125 )
2018-12-17 15:34:05 -08:00
Matthew Johnson
5589a3cc58
fix up error messages
2018-12-17 15:13:41 -08:00
Peter Hawkins
f78bd8a550
Merge pull request #130 from hawkinsp/master
...
Progress towards calling LAPACK kernels on CPU for unbatched Cholesky and Triangular solve
2018-12-17 18:05:21 -05:00
Peter Hawkins
8c25c29295
Merge pull request #127 from j-towns/qr-jvp
...
Implement the JVP of the QR decomposition
2018-12-17 17:54:43 -05:00
Peter Hawkins
e0f421746b
Import CPU Lapack implementation conditionally to ease jaxlib upgrade.
2018-12-17 17:52:16 -05:00
Peter Hawkins
ab1ebc6bad
Add experimental warning to numpy.linalg and scipy.linalg.
2018-12-17 17:39:46 -05:00
Matthew Johnson
1f2925ea8a
add backend-specific translation table in xla.py
2018-12-17 17:30:27 -05:00
Matthew Johnson
7524f2c087
fix mean/var/std kwargs ( closes #125 )
2018-12-17 14:26:28 -08:00
Peter Hawkins
eac96ac239
Fix bugs with bool argument passing; pass PRED values as int32s instead.
2018-12-17 17:17:47 -05:00
Peter Hawkins
0333a98bab
Add triangular solve BLAS implementation.
2018-12-17 16:39:19 -05:00
Peter Hawkins
3c388b98f1
Add support for calling LAPACK primitives from SciPy from JAX linalg.
2018-12-17 16:30:27 -05:00
Roy Frostig
e0ad5bb394
add TODO comment for containers in api.make_jaxpr
2018-12-17 11:42:45 -08:00
Jamie Townsend
5cdf915fb6
Test qr jvp
2018-12-17 16:36:55 +00:00
Jamie Townsend
f5b8d97c95
Add url for qr jvp notes
2018-12-17 16:04:51 +00:00
Jamie Townsend
1743a936eb
Add qr decomposition jvp
2018-12-17 16:02:29 +00:00
Roy Frostig
f4a8e03ce1
add a basic make_jaxpr transformation to the api module
2018-12-16 13:26:02 -08:00
Roy Frostig
3b8fdb050a
wrap "jit" around generated function name
2018-12-16 13:24:20 -08:00
