ROCm/jax

Pull Requests Commits

update with command_buffer off for pallas tests

cj401-amd committed 3 days ago

8bd9eb20

Use logging instead of absl.logging in interpret/shared_memory.py

jburnim committed 3 days ago

a96db3fc

Only log elapsed time during tracing for top level jit

yashk2810 committed 3 days ago

cea418ae

Rename `manual_type=.../.manual_type/.mt` to `manual_axis_type=.../.manual_axis_type/.mat` to avoid confusion with `jnp.array(...).mT`

rdyro committed 3 days ago

6735ea09

Merge pull request #36302 from jakevdp:del-get-aval

Google-ML-Automation committed 3 days ago

8799b333

Merge pull request #36292 from jakevdp:argsort-stable

Google-ML-Automation committed 3 days ago

a1747e18

Make CPU devices report as cpu:0 instead of TFRT_CPU_0.

hawkinsp committed 3 days ago

5be30dd5

Remove deprecated symbol `jax.interpreters.batching.NotMapped`.

Jake VanderPlas committed 3 days ago

6e1c2964

Remove unused logging import in pjit.py

jburnim committed 3 days ago

8de54a1c

Remove internal get_aval alias

jakevdp committed 3 days ago

1c311816

Ignore the missing attr error for `Jaxpr._debug_info` for now.

belitskiy committed 3 days ago

5f26b11a

Removed a few unused classes

superbobry committed 3 days ago

edaec9a6

[test] remove obsolete NumPy compatibility code

jakevdp committed 3 days ago

f0508450

Merge pull request #36189 from groszewn:num_eqns

Google-ML-Automation committed 3 days ago

163aa540

Track total number of eqns contained in a jaxpr

groszewn committed 3 days ago

8fec7612

Merge pull request #36093 from jpbrodrick89:jpb/scipy-batching

Google-ML-Automation committed 3 days ago

8842cf88

Make is_equivalent_to work with AbstractMesh

yashk2810 committed 3 days ago

3f1d12f3

Fix top_k's sharding rule to error out if the input array is sharded over the axis along which top_k values are computed because that will lead to comms.

yashk2810 committed 3 days ago

521fa8b9

[Mosaic] Better type constraints for conversion ops

tlongeri committed 3 days ago

06d5f6c0

Remove deprecated jax.core.[un]mapped_aval for JAX v0.10.0.

Jake VanderPlas committed 3 days ago

db61be4c

Merge pull request #36290 from superbobry:pyrefly

Google-ML-Automation committed 3 days ago

e6074764

Simplify output memory space colors lowering.

subhankarshah committed 3 days ago

1d2b27f7

Suppressed a few pytype errors

superbobry committed 3 days ago

760d6108

Removed a few more unused functions

superbobry committed 3 days ago

1354fb67

Remove code to handle DisabledSafetyCheckKind.shape_assertions.

hawkinsp committed 3 days ago

70820be0

Removed unused `backend_compile` and the corresponding `PyClient` method

superbobry committed 3 days ago

9412c26f

Removed a few jax2tf test cases which have been skipped since 2021

superbobry committed 3 days ago

99abdd4e

[Mosaic TPU] Adjust pass signature

apaszke committed 3 days ago

01f70e5b

Skip an outdated Pallas back-compat testcase.

belitskiy committed 3 days ago

43cda87e

Merge pull request #36277 from superbobry:pyrefly

Google-ML-Automation committed 3 days ago

7d00812f

Older