ROCm/jax

Pull Requests Commits

[ROCM] add typed XLA FFI support in rocm specific code

Ruturaj4 committed 2 years ago

58b658cf

Merge pull request #22222 from mattjj:xla-computation-deprecation

a-googler committed 2 years ago

42724ebc

Merge pull request #22199 from gnecula:export_test_skip

a-googler committed 2 years ago

2e0c100f

[export] Disable serialization in export_test if flatbuffers is not installed

gnecula committed 2 years ago

cfa3c91c

Merge pull request #22181 from jakevdp:xla-abbrevs

a-googler committed 2 years ago

92ebb533

add more info to xla_computation deprecation warning

mattjj committed 2 years ago

bd166e1d

Removed inspect.Signature hackery from ``pl.BlockSpec``

superbobry committed 2 years ago

da76ebf0

Merge pull request #21982 from ayaka14732:main

a-googler committed 2 years ago

81fcc97e

Migrated ``pl.BlockSpec`` uses in JAX to the new argument order

superbobry committed 2 years ago

8bc11138

[pallas] Make pallas_test run on CPU and TPU also.

gnecula committed 2 years ago

de0fd722

Use a class with `__slots__` instead of a NamedTuple in JaxprEqn and SourceInfo, which are two tuples we build frequently.

hawkinsp committed 2 years ago

2350a73f

Small optimization when forming IR constant.

hawkinsp committed 2 years ago

063581a3

Merge pull request #22119 from dfm:cond-linear

a-googler committed 2 years ago

db13e6fc

Merge pull request #22167 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

68e96830

Merge pull request #21925 from dfm:ffi-call

a-googler committed 2 years ago

b669ab7b

Add tests for some Pallas/Mosaic issues filed recently.

Google-ML-Automation committed 2 years ago

1eaaa10a

Changed ``pl.BlockSpec`` to accept ``block_shape`` before ``index_map``

superbobry committed 2 years ago

a2a5068e

Run single controller device_put via efficient reshard if the device_set of input and the sharding is the same. The transfer_guard in the test fails before this CL.

yashk2810 committed 2 years ago

94ba6c3f

[pallas] Add limited support for shape polymorphism for TPU

gnecula committed 2 years ago

9def0f1c

Don't use the trace context in the prune_closed_jaxpr_outputs cache.

hawkinsp committed 2 years ago

55589fbf

[Pallas] Add vmap/batched key support to to_pallas_key. This is helpful for workflows where the key is split before being passed into the kernel.

justinjfu committed 2 years ago

9653f58f

Merge pull request #22132 from rajasekharporeddy:testbranch2

a-googler committed 2 years ago

7cd5b5c6

Remove `linear` parameter from `lax.cond_p`.

dfm committed 2 years ago

6becf716

Improve docs for jnp.var and std

rajasekharporeddy committed 2 years ago

6713e9dc

Merge pull request #22208 from superbobry:docs

a-googler committed 2 years ago

32a70718

Add `ffi_call` function with a similar signature to `pure_callback`.

dfm committed 2 years ago

e9b087d3

Revived the workaround for not-expanding type aliases

superbobry committed 2 years ago

e80632e6

Avoid different return value when running on single host.

andsteing committed 2 years ago

a386af44

Bump up the shard_count for GPU FFT tests

apaszke committed 2 years ago

727d1204

Update jnp.mean doc

rajasekharporeddy committed 2 years ago

9623dcfa

Older