mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-22 18:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
27,173 Commits 508 Branches 263 Tags
Commit Graph

9119 Commits

Author SHA1 Message Date
Alex Dragan
b7fee456e6 [impl] Add support for setdiff1d 2020-10-21 20:17:12 +03:00
Roman Novak
da0bff2fa8 Add lax.conv_general_dilated_patches 2020-10-20 22:58:53 -07:00
Peter Hawkins
d001ac6b8a [JAX] Add support for retaining ownership of DLPack tensors. 
Move dlpack.py contents under jax/_src/dlpack.py.

Add array interoperability test between JAX and TensorFlow using DLPack.

Fixes: https://github.com/google/jax/issues/4636
PiperOrigin-RevId: 338120910
 2020-10-20 13:07:07 -07:00
Jean-Baptiste Lespiau
114c1dd070 Also convert numpy arrays to hashable values. 
This was working for a combination of lists and tuples, but the user can also pass a numpy array. Iterable is the correct type here, not Sequence.
 2020-10-20 13:57:27 +02:00
Jake VanderPlas
c56f76db06 Add deprecation warning for indexing with non-tuple sequences 2020-10-19 14:07:40 -07:00
jax authors
e5e33f57e1 Merge pull request #4633 from hawkinsp:src 
PiperOrigin-RevId: 337681883
 2020-10-17 13:50:44 -07:00
Peter Hawkins
10b7d7d7c2 Move implementation of jax.lax into jax._src.lax. 
Remove lax_ prefixes from jax/_src/lax filenames, since they aren't needed any longer to avoid name conflicts.
 2020-10-17 16:09:21 -04: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
b2808dc8f4 Move jax.nn implementation into jax._src.nn. 2020-10-17 13:45:01 -04:00
jax authors
e9909ce008 Copybara import of the project: 
--
a396cfbbd414f6f21f0c7e8a68e6e89d202c0e84 by Peter Hawkins <phawkins@google.com>:

Move jax.nn implementation into jax._src.nn.

PiperOrigin-RevId: 337671917
 2020-10-17 10:40:21 -07:00
jax authors
15bd7f6406 Merge pull request #4628 from hawkinsp:src3 
PiperOrigin-RevId: 337671365
 2020-10-17 10:31:16 -07:00
jax authors
8841dd4efd Merge pull request #4627 from hawkinsp:src2 
PiperOrigin-RevId: 337671357
 2020-10-17 10:30:56 -07:00
jax authors
ce51b976f9 Merge pull request #4626 from hawkinsp:src 
PiperOrigin-RevId: 337671310
 2020-10-17 10:27:42 -07:00
Jean-Baptiste Lespiau
1c7d726f6d Make pad more robust to non-hashable values. 
The `pad_width` can be an iterable of iterbles, thus, we need to cast both the inner and outer list to tuples.
Also, the values should be intergers, but some use-cases (mostly Tensorflow) pass e.g. DeviceArray(1) as values in there.
 2020-10-17 18:22:45 +02:00
Peter Hawkins
ef57858deb Move jax.ops implementation into jax._src.ops. 2020-10-17 11:45:28 -04:00
Peter Hawkins
c876efcc2c Move implementation of jax.image to jax._src.image. 2020-10-17 11:37:55 -04:00
Peter Hawkins
a396cfbbd4 Move jax.nn implementation into jax._src.nn. 2020-10-17 11:31:19 -04:00
Peter Hawkins
aa107cf1f4 Move jax.numpy internals into jax._src.numpy. 2020-10-16 20:35:19 -04:00
Peter Hawkins
6acb46516e Move most of the implementation of jax.scipy into jax._src.scipy. 2020-10-16 17:04:25 -04:00