mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
18,992 Commits 501 Branches 263 Tags
Commit Graph

18992 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Peter Hawkins
c4368351d2 Add support for bool dlpack values. 
PiperOrigin-RevId: 599199196
 2024-01-17 09:30:42 -08:00
jax authors
1bd22b0fe1 Merge pull request #19388 from jakevdp:rng-bits-test 
PiperOrigin-RevId: 599131268
 2024-01-17 04:38:59 -08:00
jax authors
7520363c04 Update XLA dependency to use revision 
c8adbc4bfb.

PiperOrigin-RevId: 599021174
 2024-01-16 18:52:45 -08:00
jax authors
322606978f Merge pull request #19389 from jakevdp:logsumexp 
PiperOrigin-RevId: 599002619
 2024-01-16 17:07:11 -08:00
Jake VanderPlas
7d6a134f4e logsumexp: use NumPy 2.0 convention for complex sign 2024-01-16 16:15:06 -08:00
jax authors
08837a9919 Merge pull request #19387 from jakevdp:searchsorted-test 
PiperOrigin-RevId: 598935247
 2024-01-16 13:22:38 -08:00
Jake VanderPlas
b0ed801661 random_test: test random.bits directly 2024-01-16 13:19:29 -08:00
jax authors
7f661deca3 Merge pull request #19385 from jakevdp:logistic-test 
PiperOrigin-RevId: 598935013
 2024-01-16 13:14:15 -08:00
jax authors
e20de94c5d Merge pull request #19320 from jakevdp:unique-inverse 
PiperOrigin-RevId: 598924975
 2024-01-16 12:42:11 -08:00
Jake VanderPlas
5bad0dbddd TST: make searchsorted test compatible with NumPy 2.0 2024-01-16 12:14:59 -08:00
jax authors
75b6d98367 Merge pull request #19351 from IlayMenahem:cdf-ppf 
PiperOrigin-RevId: 598914990
 2024-01-16 12:08:36 -08:00
Jake VanderPlas
ae497e8b3f test: fix dtype in test of lax.logistic 2024-01-16 11:59:30 -08:00
Jake VanderPlas
fa6d3f26ff jnp.unique: make return_inverse shape match NumPy 2.0 2024-01-16 11:47:45 -08:00
jax authors
a20fc45d65 Merge pull request #19384 from jakevdp:keyarray-doc 
PiperOrigin-RevId: 598906801
 2024-01-16 11:43:08 -08:00
ilay menahem
390e90361a Add .hypothesis/ directory to .gitignore 
and ppf and cdf to scipy.stats.uniform
 2024-01-16 18:59:52 +00:00
jax authors
c0d51e7dde Merge pull request #19381 from jakevdp:fix-diff 
PiperOrigin-RevId: 598884533
 2024-01-16 10:35:05 -08:00
jax authors
94b2da6a3b Merge pull request #19302 from carlosgmartin:scipy-stats-sem 
PiperOrigin-RevId: 598884144
 2024-01-16 10:34:45 -08:00
Jake VanderPlas
1947104212 JEP 9263: update deprecation status of KeyArray and PRNGKeyArray 2024-01-16 10:33:14 -08:00
jax authors
7307e9024b Merge pull request #19380 from jakevdp:any-all-where 
PiperOrigin-RevId: 598880307
 2024-01-16 10:26:16 -08:00
jax authors
c21602ef88 Merge pull request #19367 from google:dependabot/github_actions/actions/cache-3.3.3 
PiperOrigin-RevId: 598876775
 2024-01-16 10:17:45 -08:00
jax authors
7b0c8e4c83 Merge pull request #19242 from yixiaoer:fix-typo0 
PiperOrigin-RevId: 598876766
 2024-01-16 10:08:53 -08:00
jax authors
abe820c1e8 Merge pull request #19377 from superbobry:main 
PiperOrigin-RevId: 598866324
 2024-01-16 09:33:29 -08:00
Jake VanderPlas
17f5658db8 jnp.diff: support scalar prepend/append 2024-01-16 08:46:44 -08:00
jax authors
620fe815a4 Merge pull request #19366 from google:dependabot/github_actions/actions/upload-artifact-4.1.0 
PiperOrigin-RevId: 598850732
 2024-01-16 08:35:13 -08:00
Jake VanderPlas
b14552e811 [typing] add where paramter in any/all type declaration 2024-01-16 08:29:14 -08:00
jax authors
64f40e1201 Merge pull request #19375 from mattjj:approx-top-k-rep-rule 
PiperOrigin-RevId: 598839094
 2024-01-16 07:47:27 -08:00
Sergei Lebedev
1e9f96a574 Include Triton files into the jaxlib wheel 
This PR is based on #19368.
 2024-01-16 15:28:12 +00:00
Sergei Lebedev
af49b01e1f Migrated a subset of triton.compat to directly use IR builders 
PiperOrigin-RevId: 598826331
 2024-01-16 06:46:08 -08:00
Matthew Johnson
30c0fc4c5f [shard-map] add approx_top_k replication rule 2024-01-16 03:59:40 -08:00
jax authors
ab8eb896d7 Document several jax.config methods. 
PiperOrigin-RevId: 598775913
 2024-01-16 02:26:41 -08:00
Sharad Vikram
ef8069abaf [Pallas/Mosaic] Add support for overriding VMEM limit for kernels 
PiperOrigin-RevId: 598702707
 2024-01-15 19:52:31 -08:00
jax authors
3a4bf147d6 Update XLA dependency to use revision 
ea6abf73eb.

PiperOrigin-RevId: 598701025
 2024-01-15 19:39:22 -08:00
jax authors
55c9ed614c [splash_attention] Fix empty mask corner cases in mask_info construction 
* When a mask is completely zero'd out, then return trivial data_next and mask_next tensors
* When shrinking the launch grid, handle the case in which all the blocks in a row are masked-out. These rows will be completely padded.

PiperOrigin-RevId: 598665385
 2024-01-15 14:41:32 -08:00
jax authors
424cece35a Merge pull request #19363 from gnecula:poly_improve_max 
PiperOrigin-RevId: 598639618
 2024-01-15 11:09:02 -08:00
dependabot[bot]
2d8760f002
Bump actions/cache from 3.3.2 to 3.3.3 
Bumps [actions/cache](https://github.com/actions/cache) from 3.3.2 to 3.3.3.
- [Release notes](https://github.com/actions/cache/releases)
- [Changelog](https://github.com/actions/cache/blob/main/RELEASES.md)
- [Commits](704facf57e...e12d46a63a)

---
updated-dependencies:
- dependency-name: actions/cache
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
 2024-01-15 18:02:31 +00:00
dependabot[bot]
3facc199ed
Bump actions/upload-artifact from 4.0.0 to 4.1.0 
Bumps [actions/upload-artifact](https://github.com/actions/upload-artifact) from 4.0.0 to 4.1.0.
- [Release notes](https://github.com/actions/upload-artifact/releases)
- [Commits](c7d193f32e...1eb3cb2b3e)

---
updated-dependencies:
- dependency-name: actions/upload-artifact
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
 2024-01-15 18:02:24 +00:00
jax authors
78b46043b0 Decorate jax.nn.initializers.Initializer as @typing.runtime_checkable 
Without this decorator, we get a warning from typeguard:

```
.../typeguard/_checkers.py:474: UserWarning: Typeguard cannot check the Initializer protocol because it is a non-runtime protocol. If you would like to type check this protocol, please use @typing.runtime_checkable
```

PiperOrigin-RevId: 598588778
 2024-01-15 05:44:18 -08:00
George Necula
a1286d0021 [shape_poly] Improve core.max_dim and core.min_dim 
Previously, we optimized `core.max_dim(a, b)` to `a`
if `a >= b` and to `b` if `a < b`. Now we also optimize
it to `b` if `a <= b`.

Similarly for `core.min_dim`.
At the same time we move more of the logic from `core.py`
to `shape_poly.py`.
 2024-01-15 15:10:28 +02:00
Peter Hawkins
e558feaa5e Deprecate support for the mhlo dialect. 
JAX has not used mhlo for some time, in favor of stablehlo. Deprecate support for this dialect in JAX's API and remove testing.

PiperOrigin-RevId: 598550225
 2024-01-15 02:13:40 -08:00
Peter Hawkins
912a5ef771 Disable some tests that time out or OOM under ASAN. 
PiperOrigin-RevId: 598543036
 2024-01-15 01:40:44 -08:00
jax authors
66ee75ad19 Update XLA dependency to use revision 
fbd57d0f1e.

PiperOrigin-RevId: 598476356
 2024-01-14 18:53:08 -08:00
carlosgmartin
18ecd2e4fd Add scipy.stats.sem. 2024-01-13 22:17:21 -05:00
jax authors
f539187c05 Update XLA dependency to use revision 
78ca99bf31.

PiperOrigin-RevId: 598278037
 2024-01-13 18:30:44 -08:00
Sergei Lebedev
2d8a20c413 Do not load Triton bindings eagerly in jax/lib/__init__.py 
Triton is only used by Pallas, so it makes sense to delay loading until Pallas
is imported.

PiperOrigin-RevId: 598131836
 2024-01-13 03:01:02 -08:00
George Necula
0967a797e8 [shape_poly] Protect shape_poly: rename to _shape_poly.py. 
The public APIs can be accessed through `jax.experimental.export`.

The shape_poly and serialization modules are still changing and I saw
external references to various symbols in them, even protected ones.
I have removed such references from the Google code base, and I want to take
another step to discourage direct access to its symbols.

PiperOrigin-RevId: 598119703
 2024-01-13 02:05:11 -08:00
jax authors
b8b119d9b9 Cleanup deprecated compilation cache APIs. 
Since the compilation cache is now initialized lazily,
existing APIs initialize_cache() and is_initialized()
are confusing. Deprecate these APIs.

Introduce a new API set_cache_dir() to explicitly set the
cache directory path in code.

Testing: revised unit tests, test workload.
PiperOrigin-RevId: 598073423
 2024-01-12 22:44:48 -08:00
jax authors
2fdaef4cee Update XLA dependency to use revision 
b1f2f987e5.

PiperOrigin-RevId: 598019707
 2024-01-12 18:29:07 -08:00
jax authors
7418f55987 Merge pull request #19007 from sshahrokhi:enhanced 
PiperOrigin-RevId: 598002075
 2024-01-12 17:12:16 -08:00
jax authors
1ab5a3579a Merge pull request #19348 from jakevdp:scipy-version 
PiperOrigin-RevId: 597978339
 2024-01-12 15:48:20 -08:00
Shiva Shahrokhi
65f3e4fffd making sure enhanced barrier only turns on when there is a supported TPU available. 2024-01-12 23:47:37 +00:00