ROCm/jax

Pull Requests Commits

Update test-requirements.txt

pelumi1163 committed 51 days ago

Verified e83d5294

Update pytest_rocm.yml

pelumi1163 committed 51 days ago

Verified 829d4a25

[ROCm] Add TheRock 7.14.0 to CI nightly.

mminutoli committed 51 days ago

3e02043a

[ROCm] Update the XLA pin to include customer mitigation.

mminutoli committed 51 days ago

b64b6522

Merge pull request #39678 from vfdev-5:updated-pytorch-tests

Google-ML-Automation committed 51 days ago

4d85c9d6

Fixed `test_emit_pipeline_small` on TPU v3

superbobry committed 52 days ago

616fe02f

Added torch as test-requirements.txt and updated req files

vfdev-5 committed 52 days ago

cb4cd8fe

Update XLA dependency to use revision http://github.com/openxla/xla/commit/88ca6f6e215db0b826463d44d9ce7e194fe7217a

Google-ML-Automation committed 52 days ago

c65596c0

Merge pull request #40650 from magaonka-amd:lru-cache-avoid-pwd-lookup

Google-ML-Automation committed 52 days ago

4591efc4

[pltpu.emit_pipeline] Support TransformedRefs in async prefetch.

rdyro committed 52 days ago

1194174f

[pltpu.emit_pipeline] Support async pipeline prefetch.

rdyro committed 52 days ago

a7fe545c

Merge pull request #40688 from jakevdp:bool-annotation

Google-ML-Automation committed 52 days ago

92204b9a

C++-side support for batched_copy_to_host_async.

Google-ML-Automation committed 52 days ago

2211cfe1

[typing] jax.numpy: use builtins.bool in type stubs

jakevdp committed 52 days ago

1882f6c3

Merge pull request #40685 from mattjj:push-lzwusmztputy

Google-ML-Automation committed 52 days ago

fe66f89d

[vjp3] error on duplicate backward-pass log keys

mattjj committed 52 days ago

b340b96d

Support nested fuser.fuse by simply not fusing nested fuser.fuse.

rdyro committed 52 days ago

fd123616

Merge pull request #40677 from mattjj:push-vqswxnzlktxr

Google-ML-Automation committed 52 days ago

bf06da9f

Handle nullary primitives in program_order. We do this by rewriting the nullary primitive as `jax.jit(lambda token: eqn.primitive.bind(**params))(token)` and then program_order adds the opt_barrier on the token!

yashk2810 committed 52 days ago

374f7164

implement custom_root as a HiJAX primitive with deferred derivatives

mattjj committed 52 days ago

98f33a51

[Pallas/Mosaic GPU] Add argument validation to the abstract eval rule of `set_max_registers`.

bchetioui committed 52 days ago

bed076ba

Merge pull request #40676 from mattjj:push-pmyrnmzwlqwm

Google-ML-Automation committed 52 days ago

517ad254

[hijax] preserve symbolic zeros when deriving VJPs from linearization

mattjj committed 52 days ago

d21f649a

Merge pull request #40675 from mattjj:push-ytzulzpoolxw

Google-ML-Automation committed 52 days ago

043ee43b

[hijax] support float0 in linear_transpose

mattjj committed 52 days ago

df0add27

[pallas:mgpu] Deprecated `pl.program_id` in Pallas MGPU kernels

superbobry committed 52 days ago

0c7eecfd

[Pallas:SC] Improve tests for gathers with different numbers of rows

Google-ML-Automation committed 52 days ago

e592a247

Merge pull request #40089 from jakevdp:jnp-minmax

Google-ML-Automation committed 52 days ago

2f366628

Merge pull request #35096 from carlosgmartin:variance_scaling_division_by_zero

Google-ML-Automation committed 52 days ago

9de55379

Fix num_consts calculation in fuser.fusible.

Google-ML-Automation committed 52 days ago

0896c305

Older