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

6660 Commits

Author SHA1 Message Date
jiayaobo
ae2387dc27 add random.binomial 
update

update

modify
 2023-11-19 14:51:10 +08:00
Yash Katariya
c8ef37507b Make the SpecifiedLayout class opaque. 
Also need to enabling pickling to xc.Layout so that AOT serialization continues to work.

PiperOrigin-RevId: 583684299
 2023-11-18 15:17:16 -08:00
George Necula
3601b25899 Move multi_platform_export_test.py out of jax2tf. 
This test is now independent of jax2tf. Move it out and rename it export_harnesses_multi_platform_test.py.

We disable the test in GitHub CI, because it is very large, pending
some changes to ensure it parallelizes well. The test is still
running in internal CI. This is matching the current behavior, since
jax2tf tests are only run internally.

PiperOrigin-RevId: 583603863
 2023-11-18 02:52:44 -08:00
Yash Katariya
493e2f8ae2 If a function returns no output, xla_executable.get_output_shardings() returns 1 sharding because for XLA the output is an empty tuple which has a tuple sharding. 
PiperOrigin-RevId: 583555384
 2023-11-17 20:49:03 -08:00
Peter Hawkins
8e8dc263bc Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result. 
In most cases these are more succinct.

This change does not update Pallas/Mosaic.

PiperOrigin-RevId: 583448254
 2023-11-17 11:47:14 -08:00
Peter Hawkins
30a0136813 Increase minimum jaxlib version to 0.4.19. 
0.4.19 has xla_extension version 207 and mlir_api_version 54.

PiperOrigin-RevId: 583412447
 2023-11-17 09:38:31 -08:00
jax authors
f03d937208 Merge pull request #18564 from jakevdp:array-none 
PiperOrigin-RevId: 583368265
 2023-11-17 06:33:31 -08:00
jax authors
05bba6e790 Merge pull request #18516 from gnecula:poly_non_negative 
PiperOrigin-RevId: 583303763
 2023-11-17 01:27:14 -08:00
Yash Katariya
439b89e47f Remove DefaultLayout and make None same as DefaultLayout 
PiperOrigin-RevId: 583221970
 2023-11-16 18:01:27 -08:00
Jake VanderPlas
84aa7e5c53 Deprecate passing of None to jax.numpy.array 2023-11-16 15:10:56 -08:00
jax authors
7657a0fb15 Merge pull request #18539 from NeilGirdhar:ruff 
PiperOrigin-RevId: 583105786
 2023-11-16 11:15:19 -08:00
George Necula
db29f3230e [shape_poly] Fix handling of core.non_negative_dim symbolic expressions. 
Previously we used lax.max to evaluate core.non_negative_dim, but this is
problematic if we are in a tracing context. Then, even if the operand is
a constant we produce a tracer. Change the code to check explicitly if
the operand is a constant or if it is a symbolic expression.
 2023-11-16 20:23:51 +02:00
jax authors
71a29e6e0a Merge pull request #18550 from jakevdp:in-axes-error 
PiperOrigin-RevId: 583087978
 2023-11-16 10:22:49 -08:00
Matthew Johnson
6b6b44d409 add error hint about common jnp.ones / jnp.zeros mistake 2023-11-15 19:52:16 -08:00
Neil Girdhar
3c920c0120 Switch from flake8 to Ruff 2023-11-15 22:35:52 -05:00
Matthew Johnson
2288f64563 rot90 validate argument has ndim at least 2 2023-11-15 18:24:42 -08:00
jax authors
840b5c5d6d Merge pull request #18499 from renecotyfanboy:hyp1f1_poch 
PiperOrigin-RevId: 582765493
 2023-11-15 12:25:59 -08:00
Jake VanderPlas
0bcd64ade3 jax.vmap: improve docs & error for structured in_axes 2023-11-15 11:56:53 -08:00
jax authors
946819fc0e Merge pull request #18546 from jakevdp:fix-bool-indices 
PiperOrigin-RevId: 582742255
 2023-11-15 11:22:26 -08:00
jax authors
f2c89a43dc Merge pull request #18527 from carlosgmartin:squareplus 
PiperOrigin-RevId: 582735733
 2023-11-15 11:14:13 -08:00
sdupourque
47ca51f474 implementation of poch and hyp1f1 2023-11-15 20:01:00 +01:00
George Necula
152f60d944 [export] Minor cleanup of regexp usage in export_test. 
The goal is to make the regexp more permissive, and to ensure
that upon failure the error message has enough information to
understand the fix.
 2023-11-15 18:53:29 +01:00
Yash Katariya
118d85cd6c Make the regex checking of export_tests less strict 
PiperOrigin-RevId: 582704122
 2023-11-15 09:24:59 -08:00
Jake VanderPlas
416b734567 Fix boolean indexing check with newaxis 2023-11-15 09:03:15 -08:00
Yash Katariya
5c3da219c0 Add a private API to allow setting layouts on jitted computations. 
We expose 3 modes:

* `SpecifiedLayout`: User specifies the `minor_to_major` field of the layout. Tiling not exposed yet.

* `DefaultLayout`: PJRT chooses the layout. It defaults to the current behavior.

* `AUTO`: Compiler chooses the layout. This field is not a layout per se. It's a request to get the layout from the compiler. This field cannot be on an Array or other data types. It can only be on jit.

Public API coming soon.

Co-authored-by: Roy Frostig <frostig@google.com>
PiperOrigin-RevId: 582692036
 2023-11-15 08:48:53 -08:00
carlosgmartin
9f8e1bc34a Add nn.squareplus. 2023-11-14 23:52:41 -05:00
Skye Wanderman-Milne
32a8177348 Disable failing memories_test.py on Cloud TPU 
PiperOrigin-RevId: 582444670
 2023-11-14 14:39:15 -08:00
jax authors
2bb2aa1112 Factor LIBTPU_INIT_ARGS into the compilation cache key. 
Workloads that set the environment variable LIBTPU_INIT_ARGS
expect that the cache key will be invalidated if the value
of the variable changes between runs. Today, LIBTPU_INIT_ARGS
is not used in the cache key computation. The fix is to factor
it in similar to what is done with the XLA_FLAGS environment
variable.

Testing: new unit test; test workloads.
PiperOrigin-RevId: 582423420
 2023-11-14 13:31:08 -08:00
jax authors
2356d7afd0 Merge pull request #18515 from gnecula:export_call_bool 
PiperOrigin-RevId: 582311324
 2023-11-14 07:13:43 -08:00
Peter Hawkins
95e2d3fc2b [JAX:GPU] Generalize gesvdj kernel to iterate over the unbatched Jacobi kernel in cases that we cannot use the batched kernel. 
If the gesvdj() is preferable to gesvd() absent a batch dimension, even if there is a batch dimension we should prefer a loop of gesvdj() over a loop of gesvd().

PiperOrigin-RevId: 582279549
 2023-11-14 04:52:15 -08:00
George Necula
8713b751cb [export] Fix casts to and from dynamically shaped args for booleans 
When we call an Exported module that has polymorphic shapes we need to
insert casts to and from dynamic shape args and results. This is especially
tricky for boolean arrays because in that case we use a compare op.
 2023-11-14 09:32:30 +01:00
Qiao Zhang
24b735074a Avoids some unnecessary dtype casts in Pallas attention to make fwd pass faster. Use num_warps=8 to make bwd faster. Minor fixes to make code more readable. 
Co-authored-by: Qiao Zhang <zhangqiaorjc@google.com>
Co-authored-by: Sharad Vikram <sharadmv@google.com>
PiperOrigin-RevId: 582108461
 2023-11-13 15:46:46 -08:00
Jieying Luo
21d9e50641 Disable some more profiler test for older plugins. 
PiperOrigin-RevId: 582071616
 2023-11-13 13:47:50 -08:00
Peter Hawkins
cb182b8b22 Use a Jacobi SVD solver for unbatched SVDs up to 1024x1024 on NVIDIA GPUs. 
The unbatched Jacobi solver is faster for small-moderate matrices, and the unbatched kernel doesn't have size restrictions.

Timings on T4 GPU:

Before:

------------------------------------------------------------
Benchmark                  Time             CPU   Iterations
------------------------------------------------------------
svd/m:1/n:1           263587 ns       242274 ns         2780
svd/m:2/n:1           335561 ns       298238 ns         2303
svd/m:5/n:1           337784 ns       299841 ns         2304
svd/m:10/n:1          339184 ns       300703 ns         2311
svd/m:100/n:1         359826 ns       320088 ns         2159
svd/m:500/n:1         376124 ns       338660 ns         2076
svd/m:800/n:1         375779 ns       335590 ns         2060
svd/m:1000/n:1        419171 ns       341487 ns         2072
svd/m:1/n:2           307564 ns       270663 ns         2544
svd/m:2/n:2           320928 ns       283601 ns         2487
svd/m:5/n:2           377373 ns       344228 ns         2035
svd/m:10/n:2          380557 ns       349412 ns         1953
svd/m:100/n:2         435465 ns       403496 ns         1722
svd/m:500/n:2         444610 ns       410913 ns         1680
svd/m:800/n:2         454493 ns       416495 ns         1665
svd/m:1000/n:2        492110 ns       420539 ns         1665
svd/m:1/n:5           307316 ns       275833 ns         2531
svd/m:2/n:5           374318 ns       341432 ns         2086
svd/m:5/n:5           512928 ns       470293 ns         1361
svd/m:10/n:5          589330 ns       537070 ns         1353
svd/m:100/n:5         620164 ns       580166 ns         1193
svd/m:500/n:5         636424 ns       593692 ns         1180
svd/m:800/n:5         635545 ns       595016 ns         1181
svd/m:1000/n:5        672443 ns       597387 ns         1115
svd/m:1/n:10          310013 ns       273998 ns         2520
svd/m:2/n:10          370451 ns       334489 ns         2105
svd/m:5/n:10          560037 ns       522223 ns         1274
svd/m:10/n:10         572868 ns       535388 ns         1304
svd/m:100/n:10        959802 ns       918258 ns          765
svd/m:500/n:10        955958 ns       909778 ns          758
svd/m:800/n:10        924104 ns       879512 ns          777
svd/m:1000/n:10       950140 ns       883493 ns          775
svd/m:1/n:100         351237 ns       315554 ns         2198
svd/m:2/n:100         426883 ns       390089 ns         1792
svd/m:5/n:100         601557 ns       564493 ns         1255
svd/m:10/n:100        920819 ns       880011 ns          787
svd/m:100/n:100      7902281 ns      7229220 ns           95
svd/m:500/n:100      9720727 ns      9040679 ns           79
svd/m:800/n:100      9856378 ns      8998050 ns           79
svd/m:1000/n:100     9721017 ns      9086414 ns           79
svd/m:1/n:500         371171 ns       334217 ns         2117
svd/m:2/n:500         449165 ns       411499 ns         1700
svd/m:5/n:500         620354 ns       581866 ns         1185
svd/m:10/n:500        892375 ns       847239 ns          833
svd/m:100/n:500      9564810 ns      8867540 ns           79
svd/m:500/n:500    111924035 ns    104078023 ns            7
svd/m:800/n:500    147777319 ns    142730412 ns            5
svd/m:1000/n:500   154205084 ns    149740209 ns            5
svd/m:1/n:800         372122 ns       334212 ns         2119
svd/m:2/n:800         456672 ns       419260 ns         1680
svd/m:5/n:800         691208 ns       626003 ns         1190
svd/m:10/n:800       1017694 ns       941480 ns          730
svd/m:100/n:800      9892683 ns      9091043 ns           76
svd/m:500/n:800    144134235 ns    139129722 ns            5
svd/m:800/n:800    342790246 ns    333299774 ns            2
svd/m:1000/n:800   432820082 ns    427978978 ns            2
svd/m:1/n:1000        372785 ns       335745 ns         1805
svd/m:2/n:1000        451946 ns       413341 ns         1668
svd/m:5/n:1000        618475 ns       577213 ns         1169
svd/m:10/n:1000       907729 ns       863335 ns          808
svd/m:100/n:1000     9868543 ns      9116870 ns           76
svd/m:500/n:1000   156777811 ns    152042065 ns            5
svd/m:800/n:1000   429704070 ns    424677592 ns            2
svd/m:1000/n:1000  654864311 ns    642693162 ns            1

After:
------------------------------------------------------------
Benchmark                  Time             CPU   Iterations
------------------------------------------------------------
svd/m:1/n:1           265980 ns       245433 ns         2791
svd/m:2/n:1           340203 ns       302783 ns         2288
svd/m:5/n:1           337807 ns       301916 ns         2286
svd/m:10/n:1          338064 ns       302441 ns         2297
svd/m:100/n:1         335444 ns       298440 ns         2327
svd/m:500/n:1         338025 ns       302096 ns         2272
svd/m:800/n:1         328382 ns       291740 ns         2252
svd/m:1000/n:1        397494 ns       310905 ns         2239
svd/m:1/n:2           310464 ns       274507 ns         2535
svd/m:2/n:2           319999 ns       284247 ns         2515
svd/m:5/n:2           373435 ns       335919 ns         2069
svd/m:10/n:2          376327 ns       339327 ns         2056
svd/m:100/n:2         385061 ns       349258 ns         2003
svd/m:500/n:2         392352 ns       355735 ns         1932
svd/m:800/n:2         410736 ns       370677 ns         1881
svd/m:1000/n:2        494326 ns       405603 ns         1721
svd/m:1/n:5           316735 ns       277292 ns         2538
svd/m:2/n:5           383748 ns       342218 ns         2077
svd/m:5/n:5           494204 ns       454309 ns         1476
svd/m:10/n:5          547017 ns       508184 ns         1371
svd/m:100/n:5         514537 ns       476761 ns         1460
svd/m:500/n:5         544656 ns       504877 ns         1381
svd/m:800/n:5         642590 ns       599314 ns         1159
svd/m:1000/n:5        706166 ns       621209 ns         1106
svd/m:1/n:10          310825 ns       274374 ns         2511
svd/m:2/n:10          381316 ns       344202 ns         2094
svd/m:5/n:10          565469 ns       526759 ns         1266
svd/m:10/n:10         576111 ns       537286 ns         1299
svd/m:100/n:10        653250 ns       613392 ns         1137
svd/m:500/n:10        690532 ns       645828 ns         1080
svd/m:800/n:10        763924 ns       723677 ns          959
svd/m:1000/n:10       940342 ns       855517 ns          818
svd/m:1/n:100         306134 ns       271533 ns         2526
svd/m:2/n:100         374680 ns       339298 ns         2071
svd/m:5/n:100         576926 ns       539062 ns         1228
svd/m:10/n:100        656806 ns       615171 ns         1123
svd/m:100/n:100      3295164 ns      3138621 ns          223
svd/m:500/n:100      4269347 ns      4166000 ns          168
svd/m:800/n:100      4656541 ns      4522247 ns          154
svd/m:1000/n:100     6479223 ns      6354578 ns          112
svd/m:1/n:500         329966 ns       289083 ns         2440
svd/m:2/n:500         407535 ns       366794 ns         1947
svd/m:5/n:500         567367 ns       522809 ns         1336
svd/m:10/n:500        712307 ns       657608 ns         1065
svd/m:100/n:500      4262986 ns      4169907 ns          167
svd/m:500/n:500     28824720 ns     28650258 ns           25
svd/m:800/n:500     29330139 ns     28677269 ns           25
svd/m:1000/n:500    30848037 ns     30089216 ns           23
svd/m:1/n:800         328620 ns       289181 ns         2329
svd/m:2/n:800         419052 ns       379483 ns         1876
svd/m:5/n:800         587366 ns       546979 ns         1269
svd/m:10/n:800        830762 ns       787923 ns          893
svd/m:100/n:800      4763633 ns      4595738 ns          152
svd/m:500/n:800     30447861 ns     29949714 ns           24
svd/m:800/n:800     94188958 ns     93488372 ns            8
svd/m:1000/n:800    94701529 ns     93394677 ns            7
svd/m:1/n:1000        351102 ns       313099 ns         2218
svd/m:2/n:1000        446543 ns       407807 ns         1708
svd/m:5/n:1000        661152 ns       616174 ns         1129
svd/m:10/n:1000       915743 ns       873397 ns          802
svd/m:100/n:1000     6434730 ns      6282779 ns          113
svd/m:500/n:1000    30244321 ns     29684290 ns           24
svd/m:800/n:1000    92727423 ns     91477078 ns            8
svd/m:1000/n:1000  169500709 ns    168358420 ns            4
PiperOrigin-RevId: 582041508
 2023-11-13 12:04:13 -08:00
jax authors
871b79925e Fix test failures when we update the abseil hashtable implementation. 
PiperOrigin-RevId: 581988519
 2023-11-13 09:24:06 -08:00
Peter Hawkins
1611e1bc41 Remove PythonJitTest from api_test.py. 
Ever since the jit-pjit merge, the "Python" jit test has actually just called the same code as the "C++" jit test. We don't have a C++-free jit path any more. Remove the "Python" tests since they don't test anything.

PiperOrigin-RevId: 581965049
 2023-11-13 08:03:23 -08:00
Junwhan Ahn
55394a0914 Roll back the optimized version of jax.block_until_ready due to test breakage 
Reverts 6cc6d093643c0265c7de4027f79879f6945e0342

PiperOrigin-RevId: 581577789
 2023-11-11 12:15:45 -08:00
jax authors
45982c8439 Update test since C PjRt API topology description is available. 
The compilation_cache_test had an exclusion since the C PjRt
topology description had not been implemented. Now that it is
available, remove the exclusion.

PiperOrigin-RevId: 581396824
 2023-11-10 16:12:06 -08:00
Jieying Luo
11236dbe34 Disable profiler test for older plugins. 
PiperOrigin-RevId: 581391435
 2023-11-10 15:58:20 -08:00
Jake VanderPlas
a9452b98a3 jnp.vectorize: support None arguments 2023-11-10 14:25:42 -08:00
Jake VanderPlas
c0f3fa00f8 [random] support key dtype in custom_jvp 
To do this, we introduce a dtype for key tangents which cannot be used
to generate random values
 2023-11-10 11:16:23 -08:00
Junwhan Ahn
6cc6d09364 Implement more efficient jax.block_until_ready(x) in C++ 
The current implementation synchronously calls `ArrayImpl.block_until_ready()` one by one. This is suboptimal when it's not cheap to query the readiness of an array. Also, calling `x.block_until_ready()` causes GIL to be acquired/released repeatedly.

To address this issue, this CL introduces a C++ implementation of `jax.block_until_ready(x)` that uses IFRT's `Array::GetReadyFuture()` to asynchronously query the readiness of all arrays and wait for them once. To preserve the previous behavior, the C++ implementation also has a slow path for any non-PyArray objects that implement `block_until_ready`.

PiperOrigin-RevId: 581302290
 2023-11-10 10:34:34 -08:00
jax authors
fc6ed3bc68 Merge pull request #18463 from jakevdp:eye-offset 
PiperOrigin-RevId: 581020737
 2023-11-09 14:11:09 -08:00
Jieying Luo
3f1900e2e3 [PJRT C API] Add a util method to get the PJRT C API version of the backend. 
Disable some memories tests which are not supported on plugin older than 0.32.

PiperOrigin-RevId: 581008059
 2023-11-09 13:30:19 -08:00
Yash Katariya
cf3c041366 Disable jax memories flag. 
PiperOrigin-RevId: 580961421
 2023-11-09 10:54:02 -08:00
Jake VanderPlas
4dd6334265 jnp.eye: handle larger-than int32 offsets 2023-11-09 10:23:49 -08:00
Sharad Vikram
8fbcfce2dd [Pallas] Enable interpreter mode as default lowering for CPU 
PiperOrigin-RevId: 580700740
 2023-11-08 16:35:31 -08:00
jax authors
62e8f4d4aa Merge pull request #18431 from jakevdp:hist-ravel 
PiperOrigin-RevId: 580597288
 2023-11-08 11:37:15 -08:00
jax authors
6efcfe8fe0 Merge pull request #18386 from shacklettbp:pallas 
PiperOrigin-RevId: 580561989
 2023-11-08 09:48:27 -08:00
Jake VanderPlas
a30d51ba2e jnp.histogram: avoid flattening input 2023-11-08 08:55:09 -08:00