ROCm/jax

Pull Requests Commits

Merge pull request #21680 from ROCm:ci_spmm

a-googler committed 1 year ago

f4c64378

Merge pull request #21733 from dfm:ffi-capsule-docstring

a-googler committed 1 year ago

af90464b

Merge pull request #21690 from rajasekharporeddy:testbranch1

a-googler committed 1 year ago

bd499a92

Merge pull request #21501 from jakevdp:softmax-inf-doc

a-googler committed 1 year ago

98d7235a

Merge pull request #21731 from tttc3:cross-product-typo

a-googler committed 1 year ago

1459ac04

Merge pull request #21692 from rajasekharporeddy:testbranch2

a-googler committed 1 year ago

2899c9fa

Merge pull request #21656 from yamlyeti:yamlyeti-patch-1

a-googler committed 1 year ago

30feb352

Edit `pycapsule` docstring to provide a little bit more context

dfm committed 1 year ago

1fa66590

Refactor kernel function assigment

Paweł Paruzel committed 1 year ago

5fcd50b7

fp8 matmul in pallas

Google-ML-Automation committed 1 year ago

f51af87f

[export] Remove old deprecated APIs for jax.experimental.export.

gnecula committed 1 year ago

3914cb41

fix typo in `jax.numpy.linalg.cross` docstring

tttc3 committed 1 year ago

21f71c6b

Added debug_callback to the list of exclusions in jax2tf/tests/primitives_test.py

superbobry committed 1 year ago

5d6413ce

Add missing arguments for jnp.extract's python binding signature.

Google-ML-Automation committed 1 year ago

c01c9840

Improve docs for jnp.angle and jnp.flip

rajasekharporeddy committed 1 year ago

6d94ae32

Improve documentation for jnp.fliplr and jnp.flipud

rajasekharporeddy committed 1 year ago

6d85c389

Merge pull request #21710 from jakevdp:fix-jax2tf

a-googler committed 1 year ago

625ea07a

rename `Specialized` to `Traced` (and `specialize` to `trace`)

froystig committed 1 year ago

ea6dfd19

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

dd40d885

pallas/mosaic test: avoid leaking global config state

jakevdp committed 1 year ago

a2c31f4d

Merge pull request #21711 from jakevdp:setup-module

a-googler committed 1 year ago

a1b58604

test cleanup: use ExitStack to reduce test boilerplate

jakevdp committed 1 year ago

a861c55a

Merge pull request #21716 from gnecula:exp_rename_sharding

a-googler committed 1 year ago

d457f9a1

[export] Rename in_shardings and out_shardings fields.

gnecula committed 1 year ago

01ee768f

Implement `lower` in terms of `specialize`

yashk2810 committed 1 year ago

aee62e48

Merge pull request #21484 from dfm:custom-call-lowering

a-googler committed 1 year ago

90c83bb1

Reverts dfe61285093ff826e1ad23bb36b77a42c01040b4

marksandler2 committed 1 year ago

2c246df4

Remove `jaxpr` and `name` from `Lowered` because `specialize` already has those. This keeps the abstraction boundary clear. Adapt `export` to use `specialize`.

yashk2810 committed 1 year ago

fbf2a62a

[Mosaic] Expand vector.shape_cast support for sublane (un)folding no-ops

tlongeri committed 1 year ago

a65d3ae0

jax2tf_test: ensure no modification of global config

jakevdp committed 1 year ago

48355cde

Newer Older