ROCm/jax

Pull Requests Commits

bake sane defaults into build.py

JehandadKhan committed 109 days ago

5b040ae6

fix ROCm build.py with sane defaults

JehandadKhan committed 140 days ago

d9c35081

Update XLA dependency to use revision http://github.com/openxla/xla/commit/3518079b32330932232e4c83da0e59f8e371bcb8

Google-ML-Automation committed 148 days ago

54e9c654

Reverts 6a4f44085525c2f72ef291f19453b2e000a299a8

hawkinsp committed 148 days ago

e5892751

epiqueras committed 148 days ago

0af318d7

epiqueras committed 148 days ago

a4be7b76

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

Google-ML-Automation committed 149 days ago

99c31aee

Deprecate jax.experimental.pallas.{enable_debug_checks,debug_checks_enabled}

gnecula committed 149 days ago

a740d897

Remove obsolete libtpu version checks in TPU Pallas MPMD tests.

IvyZX committed 149 days ago

4014fed4

[Pallas Fuser] Add support for dot_general

sharadmv committed 149 days ago

f6dfe6df

Add philox2x32 PRNG impl

Jake VanderPlas committed 149 days ago

dfc754d1

Expose `flatten_one_level` to `jax.tree_util`.

IvyZX committed 149 days ago

991b35ab

Reverts 0ec534d974f320ea466540db755998af72b857b5

Jake VanderPlas committed 149 days ago

be2d5fd4

[Mosaic] Add tpu.core_id and tpu.subcore_id ops

vsytch committed 149 days ago

58449937

Add experimental jax control dependency API.

mwhittaker committed 149 days ago

61acd9ac

Add a JAX/Pallas flag to poison allocated buffers in Mosaic TPU lowering.

sharadmv committed 149 days ago

d24ccb6a

jax: replace trace_to_jaxpr_dynamic with trace_to_jaxpr in state and interpreters

Google-ML-Automation committed 149 days ago

d3957b12

Small optimizations to stablehlo lowering:

hawkinsp committed 149 days ago

6a4f4408

[pallas:sc] Defaulted `needs_layout_passes` to True

superbobry committed 149 days ago

e4f9ae0c

[pallas] Moved the configuration values to jax/_src/config.py

gnecula committed 149 days ago

9d036147

[Pallas] Enable disabled TPU Interpret Mode test of example kernel

jburnim committed 149 days ago

8c38c327

[pallas:sc] `pl.kernel` is now the only API for writing Pallas SC kernels

superbobry committed 149 days ago

6820cf61

[state] Fixed a latent bug in `_call_primitive_discharge_rule`

superbobry committed 149 days ago

456b22d5

Update XLA dependency to use revision http://github.com/openxla/xla/commit/9b635916ecc6df6efee62d8e4b0c7ef87ef84d69

belitskiy committed 149 days ago

d154b809

Update XLA dependency to use revision http://github.com/openxla/xla/commit/31d42c4b1652eb6172577eeade3f64e4edc4d87d

Google-ML-Automation committed 150 days ago

11af0680

Improve compilation throughput in multi-threaded auto-tuning environments. This CL updates the `GetModuleImage` and the `GetFunctionForContext` to shorten the scope of mutexes, use reader locks wherever applicable and reduce the number of lookups.

Google-ML-Automation committed 150 days ago

1ecdea11

Make JaxprEqnContext a bit easier to work with making it easier to add new configs to it.

yashk2810 committed 150 days ago

46f75a2a

[Pallas MPMD] Refactor tracing context

sharadmv committed 150 days ago

7c01f9b1

jax: replace trace_to_jaxpr_dynamic with trace_to_jaxpr in state and interpreters

dougalm committed 150 days ago

dfb842de

add default no-unwrap implementation of Tracer.full_lower

rdyro committed 150 days ago

e6be71e2

Older