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

6198 Commits

Author SHA1 Message Date
Jake VanderPlas
814b32a44b tree_all: add support for is_leaf 2024-06-10 09:46:15 -07:00
Sergei Lebedev
5e7ad600e2 Removed the double re-exporting of Pallas GPU/TPU APIs 
jax.experimental.pallas.{gpu,tpu} now import directly from the relevant
jax._src.pallas.{triton,mosaic} submodules.

PiperOrigin-RevId: 641875127
 2024-06-10 05:59:09 -07:00
George Necula
2ade7e7526 [pallas] Move the hardware_generation query in the code path that needs it 
This change allows us to lower and export Pallas calls even
on machines that do not have TPUs, in many cases.

PiperOrigin-RevId: 641841079
 2024-06-10 03:13:36 -07:00
jax authors
af95803d00 Merge pull request #21759 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 641831969
 2024-06-10 02:29:12 -07:00
rajasekharporeddy
775c6f8727 Fix Typos in docs and one error message 2024-06-10 11:38:01 +05:30
George Necula
14d87d3bf7 [export] Move the export implementation to jax._src.export. 
This is part of the work to move the export APIs out
of jax.experimental. For now, the way to use this
implementation is still through `jax.experimental.export`.

Had to add a few "#type ignore" to the _export.py because
previously the file was exempt from internal pytype.
Will try to fix these in a later PR.

PiperOrigin-RevId: 641688200
 2024-06-09 08:59:50 -07:00
jax authors
b486a95186 Merge pull request #21507 from renecotyfanboy:main 
PiperOrigin-RevId: 641429523
 2024-06-07 20:28:23 -07:00
jax authors
d32404020b Avoid "min() arg is an empty sequence" error after enabling "jax_explain_cache_misses". 
PiperOrigin-RevId: 641381432
 2024-06-07 15:52:35 -07:00
sdupourque
751d59ce67 increase default precision for hyp1f1 2024-06-08 00:38:51 +02:00
Yash Katariya
44a13c9d4b Merge code between make_jaxpr and jit(f).trace. 
The semantics of `make_jaxpr` are preserved here i.e. `make_jaxpr` still closes over tracers but `jit(f).trace` doesn't.

Since we can keep the existing behavior and still merge the implementation is a good cleanup!

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

PiperOrigin-RevId: 641347140
 2024-06-07 13:48:31 -07:00
jax authors
25cc84b879 Merge pull request #21615 from selamw1:append_doc 
PiperOrigin-RevId: 641344856
 2024-06-07 13:39:57 -07:00
jax authors
dfc6076db2 Merge pull request #21744 from superbobry:typing 
PiperOrigin-RevId: 641339815
 2024-06-07 13:23:31 -07:00
jax authors
7d913f763a Merge pull request #21298 from oliverdutton:pallas_interpreter_indexing_fix 
PiperOrigin-RevId: 641325047
 2024-06-07 12:29:31 -07:00
Sergei Lebedev
0786da8fd8 Removed unnecessary mypy exclusions from pyproject.toml 
* 2/3 files type check just fine now
* the remaining one could be handled via a file-level directive
 2024-06-07 20:07:42 +01:00
jax authors
af90464b53 Merge pull request #21733 from dfm:ffi-capsule-docstring 
PiperOrigin-RevId: 641307843
 2024-06-07 11:27:41 -07:00
jax authors
bd499a921e Merge pull request #21690 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 641292860
 2024-06-07 10:38:07 -07:00
jax authors
98d7235aee Merge pull request #21501 from jakevdp:softmax-inf-doc 
PiperOrigin-RevId: 641291919
 2024-06-07 10:34:40 -07:00
jax authors
1459ac04a8 Merge pull request #21731 from tttc3:cross-product-typo 
PiperOrigin-RevId: 641285460
 2024-06-07 10:18:35 -07:00
jax authors
2899c9fada Merge pull request #21692 from rajasekharporeddy:testbranch2 
PiperOrigin-RevId: 641285369
 2024-06-07 10:15:22 -07:00
Dan Foreman-Mackey
1fa66590d1 Edit pycapsule docstring to provide a little bit more context 
The docstring for the recently added `pycapsule` function in
`jax.extend.ffi` didn't conform to our usual docstring format, so I
updated it and added a little bit more context.
 2024-06-07 13:07:03 -04:00
jax authors
f51af87fc5 fp8 matmul in pallas 
PiperOrigin-RevId: 641254832
 2024-06-07 08:17:06 -07:00
George Necula
3914cb415d [export] Remove old deprecated APIs for jax.experimental.export. 
See CHANGELOG.md.
The deprecation period has passed.

Also replace deprecated .call_exported with .call in tests.

PiperOrigin-RevId: 641236222
 2024-06-07 06:52:10 -07:00
tttc3
21f71c6b66 fix typo in jax.numpy.linalg.cross docstring 2024-06-07 13:43:51 +01:00
rajasekharporeddy
6d94ae3274 Improve docs for jnp.angle and jnp.flip 2024-06-07 10:03:07 +05:30
rajasekharporeddy
6d85c3890d Improve documentation for jnp.fliplr and jnp.flipud 2024-06-07 09:58:02 +05:30
Roy Frostig
ea6dfd1947 rename Specialized to Traced (and specialize to trace) 
PiperOrigin-RevId: 641076488
 2024-06-06 17:43:08 -07:00
jax authors
a1b5860427 Merge pull request #21711 from jakevdp:setup-module 
PiperOrigin-RevId: 641049524
 2024-06-06 15:59:07 -07:00
Jake VanderPlas
a861c55a28 test cleanup: use ExitStack to reduce test boilerplate 2024-06-06 14:18:27 -07:00
jax authors
d457f9a116 Merge pull request #21716 from gnecula:exp_rename_sharding 
PiperOrigin-RevId: 641017765
 2024-06-06 14:17:10 -07:00
George Necula
01ee768f73 [export] Rename in_shardings and out_shardings fields. 
We rename `in_shardings` to `in_shardings_hlo` to remove confusion
with JAX's use of `in_shardings`.
We also rename `xla_compatible_in_sharding` to `in_shardings_jax`
since we do not have a XLACompatibleSharding type anymore.
 2024-06-06 22:00:16 +01:00
Yash Katariya
aee62e4874 Implement lower in terms of specialize 
PiperOrigin-RevId: 641005643
 2024-06-06 13:39:07 -07:00
jax authors
90c83bb1e2 Merge pull request #21484 from dfm:custom-call-lowering 
PiperOrigin-RevId: 640996459
 2024-06-06 13:10:28 -07:00
Mark Sandler
2c246df439 Reverts dfe61285093ff826e1ad23bb36b77a42c01040b4 
PiperOrigin-RevId: 640987745
 2024-06-06 12:41:17 -07:00
Yash Katariya
fbf2a62aa1 Remove jaxpr and name from Lowered because specialize already has those. This keeps the abstraction boundary clear. Adapt export to use specialize. 
PiperOrigin-RevId: 640968129
 2024-06-06 11:38:56 -07:00
jax authors
82516c5d4f Merge pull request #21694 from rajasekharporeddy:doc_typos 
PiperOrigin-RevId: 640956334
 2024-06-06 11:05:37 -07:00
jax authors
15e41a620f Merge pull request #21702 from hawkinsp:cudnnplug 
PiperOrigin-RevId: 640932820
 2024-06-06 09:58:12 -07:00
Peter Hawkins
971ab0fba2 Make CuDNN SDPA API work with JAX with a CUDA plugin configuration. 2024-06-06 12:09:19 -04:00
Christos Perivolaropoulos
18e55d567f [test_utils] Fix the encoding of capture_stdout so it works on windows. 
PiperOrigin-RevId: 640910749
 2024-06-06 08:43:25 -07:00
Dan Foreman-Mackey
ac560c0d90 Add helper function for building custom call lowering rules 
This function provides sensible defaults for custom call lowering rules
with the goal of reducing the amount of boilerplate required for
implementing custom calls.

Co-authored-by: Sergei Lebedev <slebedev@google.com>
 2024-06-06 11:34:08 -04:00
jax authors
fe9c1606fc Merge pull request #21655 from gnecula:exp_lower 
PiperOrigin-RevId: 640898362
 2024-06-06 07:58:41 -07:00
Peter Hawkins
dfe6128509 Reverts da816d34eaad6a1c6536959ccb4bfee4466c037d 
PiperOrigin-RevId: 640886105
 2024-06-06 07:10:09 -07:00
rajasekharporeddy
aee310b54d Fix some doc typos 2024-06-06 15:19:53 +05:30
George Necula
079eea5669 [export] Add a LoweringParameters.for_export boolean context for exporting 
This boolean context field is set only when we are lowering for
exporting. It can be used, e.g., to adapt the lowering rules
for the export case.
 2024-06-06 06:28:58 +01:00
Yash Katariya
55d0f5ef8f Add lower to specialize making it a true Stage. 
So now users can do:

```
specialized = jax.jit(f).specialize(*args)
print(specialized.jaxpr, specialized.out_info)

lowered = specialized.lower()

compiled = lowered.compile()
```
PiperOrigin-RevId: 640737396
 2024-06-05 19:54:41 -07:00
jax authors
d117305dba Merge pull request #21685 from pkgoogle:better_bitwise_count_doc 
PiperOrigin-RevId: 640717333
 2024-06-05 18:29:20 -07:00
Yash Katariya
228adb4a4a Add specialize on jax.jit and make it a Stage. 
Eventually, we should use this in jax.make_jaxpr and delete all the duplicated code.

PiperOrigin-RevId: 640707223
 2024-06-05 17:46:14 -07:00
jax authors
2a08b96306 Merge pull request #21629 from ayaka14732:typo1 
PiperOrigin-RevId: 640703754
 2024-06-05 17:31:57 -07:00
Mark Sandler
da816d34ea Makes global_shape optional for jax.make_array_from_process_local_data. 
PiperOrigin-RevId: 640695090
 2024-06-05 16:58:08 -07:00
Piseth Ky
61044cee49 better docs for bitwise_count 2024-06-05 16:16:54 -07:00
jax authors
bf59a67bf0 Merge pull request #21657 from rajasekharporeddy:test_branch7 
PiperOrigin-RevId: 640652642
 2024-06-05 14:35:35 -07:00