ROCm/jax

Pull Requests Commits

Add TODO responsibility

pelumi1163 committed 10 hours ago

db039b81

Disable torch for py3.15

pelumi1163 committed 11 hours ago

4571c720

Merge pull request #40643 from jakevdp:acosh-grad

Google-ML-Automation committed 11 hours ago

7a67070f

[Pallas/Mosaic GPU] Make `Barrier` and `ClusterBarrier`s default `num_barriers` 0D.

bchetioui committed 11 hours ago

1ab54cd8

Merge pull request #40509 from PhilHippo:docs-pytree-aux-data

Google-ML-Automation committed 12 hours ago

d42ad910

[autodiff] acosh: avoid overflow for gradient at large inputs

jakevdp committed 12 hours ago

35683e94

Resolve hardcoded TPU tiling shapes in JAX Pallas.

Google-ML-Automation committed 12 hours ago

fe025ad7

Show a non-None aux_data in the custom pytree examples

PhilHippo committed 18 hours ago

8f180fea

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

Google-ML-Automation committed 18 hours ago

4d85c9d6

Fixed `test_emit_pipeline_small` on TPU v3

superbobry committed 20 hours ago

616fe02f

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

vfdev-5 committed 21 hours ago

cb4cd8fe

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

Google-ML-Automation committed 23 hours ago

c65596c0

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

Google-ML-Automation committed 1 day ago

4591efc4

[pltpu.emit_pipeline] Support TransformedRefs in async prefetch.

rdyro committed 1 day ago

1194174f

[pltpu.emit_pipeline] Support async pipeline prefetch.

rdyro committed 1 day ago

a7fe545c

Merge pull request #40688 from jakevdp:bool-annotation

Google-ML-Automation committed 1 day ago

92204b9a

C++-side support for batched_copy_to_host_async.

Google-ML-Automation committed 1 day ago

2211cfe1

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

jakevdp committed 1 day ago

1882f6c3

Merge pull request #40685 from mattjj:push-lzwusmztputy

Google-ML-Automation committed 1 day ago

fe66f89d

[vjp3] error on duplicate backward-pass log keys

mattjj committed 1 day ago

b340b96d

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

rdyro committed 1 day ago

fd123616

Merge pull request #40677 from mattjj:push-vqswxnzlktxr

Google-ML-Automation committed 1 day 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 1 day ago

374f7164

implement custom_root as a HiJAX primitive with deferred derivatives

mattjj committed 1 day ago

98f33a51

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

bchetioui committed 1 day ago

bed076ba

Merge pull request #40676 from mattjj:push-pmyrnmzwlqwm

Google-ML-Automation committed 1 day ago

517ad254

[hijax] preserve symbolic zeros when deriving VJPs from linearization

mattjj committed 1 day ago

d21f649a

Merge pull request #40675 from mattjj:push-ytzulzpoolxw

Google-ML-Automation committed 1 day ago

043ee43b

[hijax] support float0 in linear_transpose

mattjj committed 1 day ago

df0add27

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

superbobry committed 1 day ago

0c7eecfd

Older