ROCm/jax

Pull Requests Commits

Fix float16 tolerance in testReducerInitial

magaonka-amd committed 12 days ago

cc61da68

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

Google-ML-Automation committed 12 days ago

4d85c9d6

Fixed `test_emit_pipeline_small` on TPU v3

superbobry committed 12 days ago

616fe02f

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

vfdev-5 committed 12 days ago

cb4cd8fe

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

Google-ML-Automation committed 12 days ago

c65596c0

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

Google-ML-Automation committed 12 days ago

4591efc4

[pltpu.emit_pipeline] Support TransformedRefs in async prefetch.

rdyro committed 12 days ago

1194174f

[pltpu.emit_pipeline] Support async pipeline prefetch.

rdyro committed 12 days ago

a7fe545c

Merge pull request #40688 from jakevdp:bool-annotation

Google-ML-Automation committed 12 days ago

92204b9a

C++-side support for batched_copy_to_host_async.

Google-ML-Automation committed 12 days ago

2211cfe1

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

jakevdp committed 12 days ago

1882f6c3

Merge pull request #40685 from mattjj:push-lzwusmztputy

Google-ML-Automation committed 12 days ago

fe66f89d

[vjp3] error on duplicate backward-pass log keys

mattjj committed 12 days ago

b340b96d

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

rdyro committed 12 days ago

fd123616

Merge pull request #40677 from mattjj:push-vqswxnzlktxr

Google-ML-Automation committed 12 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 12 days ago

374f7164

implement custom_root as a HiJAX primitive with deferred derivatives

mattjj committed 12 days ago

98f33a51

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

bchetioui committed 12 days ago

bed076ba

Merge pull request #40676 from mattjj:push-pmyrnmzwlqwm

Google-ML-Automation committed 12 days ago

517ad254

[hijax] preserve symbolic zeros when deriving VJPs from linearization

mattjj committed 12 days ago

d21f649a

Merge pull request #40675 from mattjj:push-ytzulzpoolxw

Google-ML-Automation committed 12 days ago

043ee43b

[hijax] support float0 in linear_transpose

mattjj committed 13 days ago

df0add27

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

superbobry committed 13 days ago

0c7eecfd

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

Google-ML-Automation committed 13 days ago

e592a247

Merge pull request #40089 from jakevdp:jnp-minmax

Google-ML-Automation committed 13 days ago

2f366628

Merge pull request #35096 from carlosgmartin:variance_scaling_division_by_zero

Google-ML-Automation committed 13 days ago

9de55379

Fix num_consts calculation in fuser.fusible.

Google-ML-Automation committed 13 days ago

0896c305

Initialize `abstract_mesh_context_manager` with `empty_abstract_mesh` and `layout_tracing_mode` with `LayoutMode.Auto` and `device_context` with `empty_concrete_mesh`.

yashk2810 committed 13 days ago

188e197e

Update XLA dependency to use revision http://github.com/openxla/xla/commit/047ee2ab94426a89bd6b5d73c6e7d4e616d0226d

bchetioui committed 13 days ago

7620f936

jakevdp committed 13 days ago

67892634

Older