ROCm/jax

Pull Requests Commits

implement _replace_with for PRNGKeyArray

cgarciae committed 1 year ago

85ad1fda

Simplify (and potentially accelerate) fftfreq(), as a single vectorized op instead of multiple .at[] calls.

jonbarron committed 1 year ago

8d2e9a85

Merge pull request #28130 from jakevdp:bernoulli-high

Google-ML-Automation committed 1 year ago

dd75a513

Merge pull request #28146 from vfdev-5:fix-export-mlir-module-mutation-27991

Google-ML-Automation committed 1 year ago

ff6c3706

Merge pull request #28141 from emergenz:fix-typo

Google-ML-Automation committed 1 year ago

11ae4cce

Fixed cached mlir module mutation issue in export

vfdev-5 committed 1 year ago

c4bb9b67

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b4c41355

Add support for TPU7x in Mosaic.

Google-ML-Automation committed 1 year ago

3dad9de7

emergenz committed 1 year ago

Verified ed719ecb

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

88617818

[Pallas/Fuser] Add support for pl.Element in fuser BlockSpec

Google-ML-Automation committed 1 year ago

f3224caf

fix BUILD file syntax error

rdyro committed 1 year ago

8117043d

jax.random.bernoulli: add mode='high' for improved sampling for small p

jakevdp committed 1 year ago

6c56d651

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

0573706b

Make `Shape::add_dimensions()` validate arguments by default.

Google-ML-Automation committed 1 year ago

580fe9ef

Update changelog to add information about breaking change with respect to tracing cache after sharding_in_types config was turned on which lead to `sharding` always being available on `ShapedArray`

yashk2810 committed 1 year ago

8318157f

Merge pull request #27733 from gspschmid:gschmid/jax_fwd_and_bwd

Google-ML-Automation committed 1 year ago

037dab7a

Refactor array serialization into separate JAX and tensorstore logic

rdyro committed 1 year ago

0e1b3419

Makes Effort_02 the default value for memory_fitting_level.

Google-ML-Automation committed 1 year ago

0ae613ee

Handle `sharding` param in convert_element_type's batching rule properly by adding the explicit mesh axis on dim 0

yashk2810 committed 1 year ago

80d1fbac

#sdy Do not close any partially sharded dimensions if using auto axes in a shard_map.

ZixuanJiang committed 1 year ago

d7c22eb1

Merge pull request #28063 from dfm:inline-array

Google-ML-Automation committed 1 year ago

c9fb5e57

Ensure outputs are tracers when inlining jit.

dfm committed 1 year ago

c6482ed6

Mark CliDebuggerTest as a thread-unsafe test.

hawkinsp committed 1 year ago

7f7d7805

Drop into `Auto` mode for `.at[...].set(...)` but instead of taking an `out_sharding` argument in `set`, use the input array's `sharding`. Since this is an update, after `.set`, the input array's sharding should be preserved.

yashk2810 committed 1 year ago

854b2c85

Don't use the deprecated jax.dlpack.to_dlpack in tests.

hawkinsp committed 1 year ago

934a0db4

[JAX] Remove jax.lib.xla_client.mlir_api_version and its uses.

hawkinsp committed 1 year ago

95156068

Allow the CPU collective implementation to be overridden to None.

hawkinsp committed 1 year ago

96865709

Reverts c2ba1790417ca206a4d88b25aef4d5ae510dd717

dfm committed 1 year ago

492cd3d9

Remove unused type forwarding declarations for jax.lib.xla_client.{Shape,XlaComputation}.

hawkinsp committed 1 year ago

13686bd8

Newer Older