ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a01cf9c4

[Mosaic TPU] Allow producing predicate phi node in control follow.

bythew3i committed 1 year ago

8a950b3d

Fix `JAX_ENABLE_X64` env var name for CI scripts.

Google-ML-Automation committed 1 year ago

436e522a

[XLA] Add a primitive for tagging individual XLA ops with frontend_attributes.

Google-ML-Automation committed 1 year ago

57975b18

Add type interface for `jax.nn`

Jake VanderPlas committed 1 year ago

f98ff068

layout: Implement out-of-line template method as a non-template.

pizzud committed 1 year ago

7483a3eb

[rollforward]: Add make_transfer_server_interface_factory to the py_socket_transfer jaxlib

pschuh committed 1 year ago

1dbeaf34

Finalize deprecation of `jax.core` `axis_frame`, `trace_state_clean`, & `typematch`.

Jake VanderPlas committed 1 year ago

117f2170

Merge pull request #30272 from olupton:rtx-6000-pro-test-skip

Google-ML-Automation committed 1 year ago

5fe062d7

Merge pull request #30259 from jburnim:jburnim_pallas_interpret_mode4

Google-ML-Automation committed 1 year ago

c69bde67

#sdy Make more changes to jaxlib that will be needed for supporting falling back to GSPMD in a follow up code change where XLA is updated.

bartchr808 committed 1 year ago

d73278e4

[jax:ffi] Add `skip_ffi_layout_processing` to `ffi_lowering` so current users of `mlir.custom_call` can migrate to `ffi_lowering` without passing layout information to `hlo.CustomCallOp.build_generic`.

danielsuo committed 1 year ago

b38198ef

#sdy Make changes to jaxlib that will be needed for supporting falling back to GSPMD in a follow up code change where XLA is updated.

bartchr808 committed 1 year ago

21d2433f

[Mosaic GPU] Add the `Distinct` constraint to the equation system.

bchetioui committed 1 year ago

f56bbc04

fused_attention_stablehlo_test_gpu: test skip for 12.0

olupton committed 1 year ago

3033bf2e

[rollforward]: Add transfer_server_interface to xla.cc as an extra optional factory argument

pschuh committed 1 year ago

199eb3c8

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1ebf4d56

[Pallas/TPU] Add empty_ref_like function + SMEM support for Refs

sharadmv committed 1 year ago

fcf0b728

Rollback #30206 due to downstream test failures

Google-ML-Automation committed 1 year ago

a02bba48

Use registers for state by default in emit_pipeline

Google-ML-Automation committed 1 year ago

5efa810a

[XLA] Refactor `xla_metadata.py` to resolve dependency cycle

Google-ML-Automation committed 1 year ago

04915957

Merge pull request #30206 from carlosgmartin:logmeanexp

Google-ML-Automation committed 1 year ago

0485af20

[Mosaic] Allow sublane rotation for non-sublane dim aligned shape.

WindQAQ committed 1 year ago

11e19d1b

Reverts aee89744efbd3aa93ffea7a3bab803483add9289

jakeharmon8 committed 1 year ago

4cdd05ed

Disable `too_slow` in data.draw() for test_cast_from_32bit

jakeharmon8 committed 1 year ago

772ae661

[jax:benchmark] Add tracing benchmarks for some common operations.

danielsuo committed 1 year ago

816fab4e

[Pallas] Add option to allow out-of-bound reads in TPU interpret mode.

jburnim committed 1 year ago

e322ae18

Set `check_leaks=False` in `def direct_linearize` to fix some tests

yashk2810 committed 1 year ago

4752cc94

[Mosaic] Support multiple non-contracting dims if they are collapsable.

WindQAQ committed 1 year ago

0e55e131

fix jax2tf scatter impl

yashk2810 committed 1 year ago

45f67c77

Older