ROCm/jax

Pull Requests Commits

[ROCm] Scale pytest workers by compute partition mode

magaonka-amd committed 19 days ago

ecf73c5f

[ROCm] Switch the pinned ROCm CI legs from 10.0 to 10.1

magaonka-amd committed 19 days ago

977b3395

[ROCm] Stop running the continuous wheel tests on a schedule

magaonka-amd committed 19 days ago

a16c47c5

[ROCm] Run export_harnesses_multi_platform_test_gpu downstream

magaonka-amd committed 19 days ago

8795be7f

Merge pull request #41360 from Intel-tensorflow:oneapi-contact-update

Google-ML-Automation committed 19 days ago

87d9d094

Improve tail accuracy of lax.erf_inv, erfcinv, and ndtri without FMA or float64.

hawkinsp committed 19 days ago

ffdde91a

[pallas:fuser] Use `util.foreach` instead of `safe_map` where appropriate

superbobry committed 19 days ago

83674090

[thunky] Disable hoist_constants_as_args when lowering call_jax

gnecula committed 19 days ago

6fb2b6e2

[pallas:fuser] Avoid copying the list of equations in `compute_usage`

superbobry committed 19 days ago

244ac941

[mgpu] Simplify `ptr_as_memref`.

chr1sj0nes committed 19 days ago

b665bb2b

[mgpu] Simplify `_extract_aliased_ref`.

chr1sj0nes committed 19 days ago

8d21ddc9

Merge pull request #41479 from vfdev-5:update-mgpu-ref-api-docs

Google-ML-Automation committed 19 days ago

527f6e37

Update XLA dependency to use revision http://github.com/openxla/xla/commit/919b0954c785133bf8837f1b56488ae328d00e51

Google-ML-Automation committed 20 days ago

85426483

[Mosaic:GPU] Use explicit CUDA graph node creation API with collective mosaic kernels.

PatriosTheGreat committed 20 days ago

3d6f5f57

[pallas:fuser] Treat `pl.Squeezed` like `None` in various fuser rules

superbobry committed 20 days ago

b9784683

Clean up unused forward declaration and header includes in `py_array.h`

hyeontaek committed 20 days ago

e0f38af4

Use move-only `absl::AnyInvocable` for `PyHostValue::ConvertFn`

hyeontaek committed 20 days ago

b5132e66

Merge pull request #34478 from hrideymarwah15:fix/betaln-hessian-34353

Google-ML-Automation committed 20 days ago

0905c464

Fix betaln Hessian for a == b boundary case

hrideymarwah15 committed 20 days ago

f6702c3d

Merge pull request #41440 from bhaskargurram-ai:perf/betainc-per-element-convergence

Google-ML-Automation committed 20 days ago

290228f9

[pallas:fuser] Fix scalar prefetch operand alignment in `make_kernel_function`

superbobry committed 20 days ago

34e64733

[jax] Fix precision, overflow, and nan/inf handling in jax.scipy.special.rel_entr and kl_div

Jake VanderPlas committed 20 days ago

3656fcd3

[jax] Fix vonmises.logpdf at kappa=0 and chi2.logpdf at x=0, inf for df=2

Jake VanderPlas committed 20 days ago

a8bce50c

Stop betainc's continued fraction per element once it converges

bhaskargurram-ai committed 20 days ago

141ec8e6

Add tests for AOT compilation with the persistent compilation cache on TPU.

Google-ML-Automation committed 20 days ago

00957d4d

[CI] Add `bcr.bazel.build` fallback mirrors to `bazel_downloader.cfg` and pass `bazel_dns_windows.security` in `ci/build_artifacts.sh`.

belitskiy committed 20 days ago

aadc9978

[pallas:fuser] Raise `ValueError` on `no_block_spec` outputs in `_push_block_spec_jaxpr`

superbobry committed 20 days ago

7d667ffc

Updated all references to `ClosedJaxpr` to use `Jaxpr` directly

superbobry committed 20 days ago

a01c2d84

[pallas:fuser] Populate `ctx.scalar_prefetch_fn` before calling pull rules

superbobry committed 20 days ago

1166835e

Internal change.

bchetioui committed 20 days ago

3609da50

Older