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
21,542 Commits 501 Branches 263 Tags
Commit Graph

21542 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Peter Hawkins
7f4ef63cd8 Run pyupgrade --py310-plus. 
Also apply manual fixes to import sorting and unused imports.
 2024-06-26 16:10:18 -04:00
jax authors
cdfe2df384 Merge pull request #22116 from jakevdp:matmul-precision 
PiperOrigin-RevId: 647048099
 2024-06-26 12:53:38 -07:00
jax authors
f66a14dd7f Merge pull request #22009 from pkgoogle:better_floor_divide_doc 
PiperOrigin-RevId: 647046910
 2024-06-26 12:49:07 -07:00
jax authors
62b4f9e8b8 Update XLA dependency to use revision 
c53e9f4e48.

PiperOrigin-RevId: 647029418
 2024-06-26 11:59:43 -07:00
jax authors
9842bdbc6d Merge pull request #22113 from hawkinsp:pyver 
PiperOrigin-RevId: 647028906
 2024-06-26 11:56:17 -07:00
Jake VanderPlas
bd94f2ff1c jax_default_matmul_precision: allow 'default', 'high', and 'highest' 2024-06-26 11:41:59 -07:00
jax authors
6886e286ce Merge pull request #22102 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 647018530
 2024-06-26 11:27:05 -07:00
Peter Hawkins
945fde41e4 Update minimum Python version to 3.10. 2024-06-26 13:47:14 -04:00
rajasekharporeddy
7281c39add Update jnp.min/amin docs 2024-06-26 23:12:47 +05:30
Peter Hawkins
f5f9da34e6 Update Github actions builds to use Python 3.10 instead of Python 3.9. 
Change in preparation for dropping Python 3.9 support.

PiperOrigin-RevId: 646997432
 2024-06-26 10:33:11 -07:00
jax authors
fe3c8e15a8 Merge pull request #21806 from cgarciae:cond-passthrough-outputs 
PiperOrigin-RevId: 646970169
 2024-06-26 09:13:07 -07:00
jax authors
8b22b5fbcc Merge pull request #22087 from jakevdp:validate-config 
PiperOrigin-RevId: 646963494
 2024-06-26 08:51:34 -07:00
Cristian Garcia
dae7e41ade fix cond passthrough outputs 2024-06-26 16:17:45 +01:00
jax authors
fb4ab2baa1 Merge pull request #22055 from Intel-tensorflow:yimei/remove_block_fp16_oncpu 
PiperOrigin-RevId: 646951088
 2024-06-26 08:12:42 -07:00
jax authors
b8aa4c52ee Merge pull request #22091 from abhinavgoel95:patch-2 
PiperOrigin-RevId: 646924044
 2024-06-26 06:36:25 -07:00
jax authors
6175e49f7f Merge pull request #22106 from dfm:fix-whitespace 
PiperOrigin-RevId: 646923886
 2024-06-26 06:36:09 -07:00
jax authors
93e8167e19 Merge pull request #22062 from ROCm:ci_export_test_fix 
PiperOrigin-RevId: 646923809
 2024-06-26 06:35:50 -07:00
jax authors
41d0a38fef Merge pull request #22039 from ROCm:ci_rocm_platform_fix 
PiperOrigin-RevId: 646923652
 2024-06-26 06:32:11 -07:00
jax authors
8b3b6ae883 Merge pull request #22047 from jakevdp:poly-annotations 
PiperOrigin-RevId: 646923575
 2024-06-26 06:28:00 -07:00
Dan Foreman-Mackey
1c640b25cb Fix pre-commit formatting errors 
It looks like #21916 introduced some whitespace issues that are causing
our pre-commits to fail.
 2024-06-26 06:38:51 -04:00
jax authors
96cf5d53c8 Merge pull request #21916 from ROCm:ci_pjrt 
PiperOrigin-RevId: 646793145
 2024-06-26 02:43:21 -07:00
jax authors
ca70ebbb86 Merge pull request #22099 from mattjj:fix-literal 
PiperOrigin-RevId: 646672753
 2024-06-25 18:07:41 -07:00
jax authors
af1a751f8c Merge pull request #22097 from selamw1:isreal_doc 
PiperOrigin-RevId: 646667978
 2024-06-25 17:45:40 -07:00
Matthew Johnson
ad4e5ec43c fix error message path since #22069 allowed Literals in jaxpr builder eqns 2024-06-26 00:38:23 +00:00
jax authors
a1709f65c0 Merge pull request #22098 from mattjj:remove-double-dots 
PiperOrigin-RevId: 646665167
 2024-06-25 17:36:47 -07:00
selamw1
98fca228ed isreal_docstring_added 
isreal_docstring_typo_fixed
 2024-06-25 17:13:13 -07:00
Matthew Johnson
8564b55ee2 remove double dots from an error message 2024-06-26 00:04:42 +00:00
jax authors
f03603fa02 Update XLA dependency to use revision 
9d5f98e7c0.

PiperOrigin-RevId: 646568318
 2024-06-25 12:30:31 -07:00
jax authors
e34ba4aa0d Merge pull request #21441 from vfdev-5:preserve-int-dtype-ceil-floor-trunc-ops 
PiperOrigin-RevId: 646556347
 2024-06-25 11:53:10 -07:00
jax authors
7b73e8ecf7 Merge pull request #22083 from rajasekharporeddy:testbranch2 
PiperOrigin-RevId: 646550598
 2024-06-25 11:36:26 -07:00
vfdev-5
70b4823348 Updated jnp.ceil/floor/trunc to preserve int dtypes 
Description:
- Updated jnp.ceil/floor/trunc to preserve int dtypes
- Updated tests
  - For integral dtypes but we can't yet today compare types vs numpy as numpy 2.0.0rc2 is not yet array api compliant in this case
 2024-06-25 20:26:53 +02:00
Piseth Ky
cc80f63521 better floor_division doc 
corrected equivalent operator scope, added See Also:

del whitespace

added operator equivalence note
 2024-06-25 11:08:22 -07:00
jax authors
2bdfd0cebb Merge pull request #22081 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 646539354
 2024-06-25 11:05:33 -07:00
rajasekharporeddy
ee0bf776fe Update jnp.max/amax docs 2024-06-25 23:07:59 +05:30
jax authors
b458e3cea1 Merge pull request #22082 from gnecula:pallas_doc1 
PiperOrigin-RevId: 646521859
 2024-06-25 10:18:49 -07:00
rajasekharporeddy
b975ca1932 Better doc for jnp.sum 2024-06-25 22:25:23 +05:30
Ruturaj Vaidya
385283c50b
Update build.py 2024-06-25 11:31:18 -05:00
Jake VanderPlas
fa73077146 jax.config: validate on set() 2024-06-25 09:02:32 -07:00
Abhinav Goel
7b9636f63c
Update flags gpu_performance_tips.md 2024-06-25 08:53:43 -07:00
Kevin Gleason
a4c92a454b Clean up gather/scatter StableHLO lowering. 
PiperOrigin-RevId: 646491586
 2024-06-25 08:39:50 -07:00
jax authors
50407e536e Fix a small comparison bug. This allows every last reserved byte to be used, and also avoids nonsensical error messages of the form: "Requested more bytes than we reserved space for: X > X". 
PiperOrigin-RevId: 646487396
 2024-06-25 08:25:48 -07:00
George Necula
0b0d5a08ea [pallas] Simplify the generation of the hlo.CustomCall 
We use the `mlir.custom_call` library that most other
JAX lowering rules use.

We remove the use of `hlo.TupleType` because custom calls
support multiple results. Tuples are not supported
by the `mlir.custom_call` and they complicate the logic
for input_output_aliasing and for the dynamically-shaped
results.

PiperOrigin-RevId: 646462934
 2024-06-25 06:57:03 -07:00
George Necula
8528f5127d [pallas] Break long lines in the Pallas docs 
No content changes.
 2024-06-25 13:30:17 +03:00
Junwhan Ahn
817eb7a9ee Skip broadcast_one_to_all for single-process JAX execution 
There is no need to actually perform exchange across processes if there's only one process.

PiperOrigin-RevId: 646309814
 2024-06-24 20:15:32 -07:00
Tomás Longeri
94c5d0d747 [Mosaic][apply-vector-layout] Fix possible segfault in arith.extsi/arith.extf after cl/644495447 
This only happens for layout pairs that are never inferred.

PiperOrigin-RevId: 646303509
 2024-06-24 19:54:08 -07:00
Yash Katariya
15ed2a8bcd Fix device_put of a scalar with PositionalSharding 
Fixes https://github.com/google/jax/issues/22073

PiperOrigin-RevId: 646279569
 2024-06-24 17:53:32 -07:00
Tomás Longeri
21bf3d196d [Mosaic][Python] Define __repr__ for VectorLayout 
Loosely follows the example MLIR's bindings for Attribute

PiperOrigin-RevId: 646270865
 2024-06-24 17:18:15 -07:00
jax authors
19bfaed06c Merge pull request #22000 from pkgoogle:better_copysign_doc 
PiperOrigin-RevId: 646264098
 2024-06-24 16:52:38 -07:00
Tomás Longeri
097806a033 [Mosaic][Python] Include error message in exceptions 
PiperOrigin-RevId: 646259787
 2024-06-24 16:36:26 -07:00
Piseth Ky
ee9290b9fd updating copysign docstring 
grammar fix

grammar fix
 2024-06-24 15:14:05 -07:00