mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
11,190 Commits 502 Branches 263 Tags
Commit Graph

175 Commits

Author SHA1 Message Date
Yash Katariya
8a61414a88 Delete the mesh context manager. The replacement for it is Mesh. 
PiperOrigin-RevId: 442619711
 2022-04-18 13:37:51 -07:00
Peter Hawkins
32abfa84ff Fix typo in changelog. 2022-04-18 08:17:14 -04:00
Peter Hawkins
38ea5a6bc0 Copybara import of the project: 
--
391dea76bc8fe264cf26ec93d42147f87847894d by Peter Hawkins <phawkins@google.com>:

Update version numbers after jax/jaxlib 0.3.7 release.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/10324 from hawkinsp:jaxlib 391dea76bc8fe264cf26ec93d42147f87847894d
PiperOrigin-RevId: 442311051
 2022-04-16 22:37:09 -07:00
Peter Hawkins
6b33c55450 2nd attempt at jax/jaxlib 0.3.7 release. 2022-04-15 15:20:02 -04:00
Jake VanderPlas
be5c84d409 Deprecate DeviceArray.tile method 2022-04-15 10:11:03 -07:00
Peter Hawkins
52a97f2e06 Jax 0.3.7 and jaxlib 0.3.7 release. 2022-04-15 12:02:05 -04:00
Peter Hawkins
3f1032cf33 Fix incorrect cross-reference breaking readthedocs build. 2022-04-14 16:12:48 -04:00
George Necula
d050327592 Deprecate jax.experimental.loops, step 2. 
Add deprecation warning and remove the tests.

PiperOrigin-RevId: 441828243
 2022-04-14 12:38:55 -07:00
Yash Katariya
08f28a6119 Finish jax release 
PiperOrigin-RevId: 441342919
 2022-04-12 18:13:16 -07:00
Peter Hawkins
4dc69034b0 Update version numbers after jax/jaxlib release. 2022-04-07 16:40:19 -04:00
jax authors
8b3f039252 Merge pull request #10039 from ajcr:add_scipy_linalg_rsf2csf 
PiperOrigin-RevId: 439997145
 2022-04-06 19:55:29 -07:00
Peter Hawkins
96ba290faf Jax 0.3.5 and jaxlib 0.3.5 release. 2022-04-06 23:56:41 +00:00
Alex Riley
869596fc2c Add jax.scipy.linalg.rsf2csf 2022-04-06 21:06:23 +01:00
Peter Hawkins
71a5eb263b [GPU] Force an input buffer copy for double precision complex-to-real IRFFTs. 
Fixes https://github.com/google/jax/issues/9946

PiperOrigin-RevId: 439414091
 2022-04-04 14:38:52 -07:00
Jake VanderPlas
df1ceaeeb1 Deprecate jax.tree_util.tree_multimap 2022-04-01 14:51:54 -07:00
Jake VanderPlas
b359b8ad96 Add CHANGELOG entry for #10069 2022-03-30 08:05:34 -07:00
Jake VanderPlas
093b7032a8 Implement jnp.from* array creation functions 2022-03-29 10:52:47 -07:00
Jake VanderPlas
f4d240c036 Remove lax_numpy from jax.numpy namespace 
This is a private module that was inadvertently exported in the past.
 2022-03-25 15:02:45 -07:00
dogeplusplus
7915c6ce27 Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning. 2022-03-23 20:55:22 +00:00
Jake VanderPlas
69969ef803 add random.loggamma and improve dirichlet & beta implementation 2022-03-21 08:33:11 -07:00
Matthew Johnson
4c5d8e969f update version and changelog for pypi 2022-03-18 14:16:00 -07:00
Matthew Johnson
d2b393bbf1 update version and changelog for pypi 2022-03-17 15:35:26 -07:00
Skye Wanderman-Milne
d7087abce6 Bump jax and jaxlib versions for 0.3.2 release 
Also add CPU pjit to changelog
 2022-03-16 14:31:00 -07:00
Skye Wanderman-Milne
f9775a2ced Update CHANGELOG and setup.py for jax + jaxlib 0.3.2 releases 2022-03-16 10:17:42 -07:00
jax authors
4d14899940 Add boolean flag to as_hlo_text to enable writing large constants. 
PiperOrigin-RevId: 434556535
 2022-03-14 13:46:22 -07:00
Peter Hawkins
08fbd77d90 [JAX] Deprecate .block_host_until_ready() in favor of .block_until_ready(). 
JAX kept an older name around (.block_host_until_ready()) in parallel with the new name (.block_until_ready()) to avoid breaking users. Deprecate it so we only have one name.

PiperOrigin-RevId: 433228545
 2022-03-08 09:14:40 -08:00
Jake VanderPlas
8c57ae2a19 Call _check_arraylike on inputs to broadcast_to and broadcast_arrays 2022-03-04 11:22:27 -08:00
jax authors
fb44d7c12f [JAX] Add release note for the graduration of the experimental.ann module. 
PiperOrigin-RevId: 431951602
 2022-03-02 08:58:53 -08:00
Jake VanderPlas
51727033b8 Remove duplicate changelog entry 2022-02-24 08:18:30 -08:00
Peter Hawkins
f51a05a889 Remove jax.ops.index... functions. 
These functions have been deprecated and have issued a DeprecationWarning since jax 0.2.22 in October 2021.
 2022-02-24 09:36:28 -05:00
Yash Katariya
c161c62878 Finish jax release 
PiperOrigin-RevId: 429670894
 2022-02-18 16:23:39 -08:00
Jake VanderPlas
da3aaa1960 Add deprecation warning to JaxTestCase and JaxTestLoader 2022-02-17 14:58:58 -08:00
Peter Hawkins
3e5ecfe363 Add jax.distributed and jax.dlpack to the docs. 
Reorder the doc modules into something closer to alphabetical order.

Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs.
 2022-02-17 16:10:07 -05:00
Jake VanderPlas
97512e9e44 JaxTestCase: set jax_numpy_rank_promotion='raise' by default 2022-02-14 09:22:05 -08:00
Yash Katariya
2162868ed9 Update values after release 
PiperOrigin-RevId: 427910510
 2022-02-10 20:32:53 -08:00
jax authors
5691010d2f Copybara import of the project: 
--
d42fffd849a4bac0c0c11a3346c93f07f8c64c44 by Jake VanderPlas <jakevdp@google.com>:

JaxTestCase: set numpy_rank_promotion='raise' by default
PiperOrigin-RevId: 427896974
 2022-02-10 19:08:29 -08:00
Jake VanderPlas
6324577a63 JaxTestCase: set numpy_rank_promotion='raise' by default 2022-02-10 16:54:31 -08:00
Yash Katariya
1ad3551ec9 Release jax and jaxlib 0.3.0 as per the new release process. 
PiperOrigin-RevId: 427809845
 2022-02-10 11:59:13 -08:00
Skye Wanderman-Milne
715066c624 Bump jax version to 0.2.29 and update CHANGELOG 2022-02-01 17:59:57 -08:00
Peter Hawkins
2388e353da Increase bazel version to 5.0.0 to match TensorFlow 
(8871926b0a).
 2022-01-28 21:11:02 +00:00
Peter Hawkins
be2f6a91ec Update XLA for jaxlib 0.1.76 release. 2022-01-27 14:26:16 +00:00
Peter Hawkins
74e4db47da Change the default IR dialect returned by .compiler_ir() to MHLO. 
PiperOrigin-RevId: 423091674
 2022-01-20 09:50:17 -08:00
Peter Hawkins
3fef74b2d0 [JAX] Change signature of .mhlo() method on compiler IR objects to return an ir.Module object instead of its string representation. 
It isn't free to pretty-print IR, so it's best to avoid it unless necessary. In addition, by returning an IR object, the user is now free to, say, print it with different options.

For example, one can now write things like:

```
In [1]: import numpy as np, jax, jax.numpy as jnp
In [2]: m = jax.jit(lambda x: x + jnp.array(np.arange(1000))).lower(7.).compiler_ir(dialect='mhlo')
In [3]: m.operation.print(large_elements_limit=10)
module @jit__lambda_.4 {
  func public @main(%arg0: tensor<f32>) -> tensor<1000xf32> {
    %0 = mhlo.constant opaque<"_", "0xDEADBEEF"> : tensor<1000xi32>
    %1 = "mhlo.convert"(%0) : (tensor<1000xi32>) -> tensor<1000xf32>
    %2 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1000xf32>
    %3 = mhlo.add %2, %1 : tensor<1000xf32>
    return %3 : tensor<1000xf32>
  }
}
```

Fixes https://github.com/google/jax/issues/9226

PiperOrigin-RevId: 422855649
 2022-01-19 11:04:48 -08:00
Matthew Johnson
0066533dae update version and changelog for pypi 2022-01-18 11:38:32 -08:00
jax authors
6411f8a033 Merge pull request #9184 from jakevdp:unique-nan 
PiperOrigin-RevId: 422287302
 2022-01-16 23:57:40 -08:00
jax authors
c9169fa0d5 Merge pull request #9189 from gnecula:tf_reduce_window 
PiperOrigin-RevId: 421875035
 2022-01-14 11:35:16 -08:00
Jake VanderPlas
bd157cf056 jnp.unique: properly handle NaN values 2022-01-13 15:54:07 -08:00
Jake VanderPlas
d8bdd9a19d lax.sort: regularize handling of -0.0 and -NaN 2022-01-13 13:03:41 -08:00
George Necula
5bfe1852a4 [jax2tf] Add jax2tf_associative_scan_reductions flag 
This flag allows users to match the JAX performance for
associative reductions in CPU.
See README.md for details.
 2022-01-13 15:52:18 +02:00
Matthew Johnson
1cf7d4ab5d Copybara import of the project: 
--
4fcdadbfb3f4c484fd4432203cf13b88782b9311 by Matthew Johnson <mattjj@google.com>:

add jax.ensure_compile_time_eval to public api

aka jax.core.eval_context

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/7987 from google:issue7535 4fcdadbfb3f4c484fd4432203cf13b88782b9311
PiperOrigin-RevId: 420928687
 2022-01-10 20:58:26 -08:00