ROCm/jax

Pull Requests Commits

[CI] Update CI builds to Ubuntu 22.04

MichaelHudgins committed 102 days ago

7952dbdf

[jax][host offloading] Turn on ComputeOffload tests for GPU.

basioli-k committed 102 days ago

7dc02949

implement derivatives for hyp2f1

mattbahr committed 102 days ago

c4bfc129

[pallas:sc] Added a libtpu version guard to `test_bitcast`

superbobry committed 102 days ago

6f2aa1bd

Mark make_user_context() tests as thread-unsafe.

hawkinsp committed 102 days ago

c59858f3

Fix Abseil C++ logging Python bindings.

allanrenucci committed 102 days ago

7ede5cef

[CI] Upgrade OSS TPU presubmit to v6e-8.

MichaelHudgins committed 102 days ago

5542dd79

Removed the no longer necessary jax-aval-named-shape deprecation

superbobry committed 102 days ago

fa3b7b9e

Integrate LLVM at llvm/llvm-project@b8649098a7fc

Google-ML-Automation committed 102 days ago

6a109ea3

Update XLA dependency to use revision http://github.com/openxla/xla/commit/61134bc5de9408c8248a8614c5739bd37fa1d924

Google-ML-Automation committed 102 days ago

9d430721

Add `out_sharding` parameter to `jnp.zeros_like` and `jnp.ones_like`

yashk2810 committed 102 days ago

187a6147

Add `out_sharding` argument to `jnp.asarray`.

yashk2810 committed 102 days ago

b97651f0

Reenable some GPU and TPU configurations in tests that now pass in CI.

hawkinsp committed 102 days ago

fa951172

Don't recompute avals of Literals during mlir lowering.

hawkinsp committed 102 days ago

38581b7f

Add `out_sharding` argument to `jax.numpy.array`.

yashk2810 committed 102 days ago

b86db752

Remove `try-except` blocks around the import of `_profile_data` in `jax/_src/lib/__init__.py` and `jax/_src/profiler.py`, as the module is now always expected to be present.

ZacCranko committed 102 days ago

0f45fad2

Add a temporary hack to do a trace only memory_space cast without affect the lowering of a jaxpr.

yashk2810 committed 102 days ago

b9c68582

Merge pull request #31979 from IvyZX:typo

Google-ML-Automation committed 102 days ago

2ead2dd1

Remove deprecated symbols from `jax.dlpack`, `jax.errors`, `jax.lib.xla_bridge`, `jax.lib.xla_client`, and `jax.lib.xla_extension`.

danielsuo committed 102 days ago

66760f71

Fix a typo in pipelining pseudocode

IvyZX committed 102 days ago

65f767fb

Don't recompute avals during constant folding or inlining in partial_eval.

hawkinsp committed 102 days ago

44554348

Merge pull request #31974 from jakevdp:special-impl

Google-ML-Automation committed 102 days ago

6b4df6ae

Fix TSAN builds after rules_python upgrade.

Google-ML-Automation committed 102 days ago

26e5f70d

Merge pull request #31896 from hawkinsp:install

Google-ML-Automation committed 102 days ago

7f0f7d9b

jax.scipy.special.ndtri: refactor implementation in terms of polyval

jakevdp committed 102 days ago

dec517bd

Add a variant of the test introduced in https://github.com/jax-ml/jax/pull/31990 without using a custom_vjp.

yashk2810 committed 102 days ago

372c067d

Merge pull request #32008 from jakevdp:ref-docs

Google-ML-Automation committed 102 days ago

48e73e68

[Mosaic GPU] Cast i4->bf16 through i32

apaszke committed 103 days ago

93832457

[mutable array] add missing Ref docs

jakevdp committed 103 days ago

37f021ba

Add an example to return unreduced gradients from a loop/scan over microbatches and then reduce it outside the loop.

yashk2810 committed 103 days ago

786db8d4

Newer Older