ROCm/jax

Pull Requests Commits

implement _replace_with for PRNGKeyArray

cgarciae committed 319 days ago

85ad1fda

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

jonbarron committed 319 days ago

8d2e9a85

Merge pull request #28130 from jakevdp:bernoulli-high

Google-ML-Automation committed 319 days ago

dd75a513

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

Google-ML-Automation committed 319 days ago

ff6c3706

Merge pull request #28141 from emergenz:fix-typo

Google-ML-Automation committed 319 days ago

11ae4cce

Fixed cached mlir module mutation issue in export

vfdev-5 committed 319 days ago

c4bb9b67

Update XLA dependency to use revision

Google-ML-Automation committed 319 days ago

b4c41355

Add support for TPU7x in Mosaic.

Google-ML-Automation committed 319 days ago

3dad9de7

emergenz committed 319 days ago

Verified ed719ecb

Update XLA dependency to use revision

Google-ML-Automation committed 320 days ago

88617818

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

Google-ML-Automation committed 320 days ago

f3224caf

fix BUILD file syntax error

rdyro committed 320 days ago

8117043d

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

jakevdp committed 321 days ago

6c56d651

Update XLA dependency to use revision

Google-ML-Automation committed 321 days ago

0573706b

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

Google-ML-Automation committed 321 days 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 321 days ago

8318157f

Merge pull request #27733 from gspschmid:gschmid/jax_fwd_and_bwd

Google-ML-Automation committed 321 days ago

037dab7a

Refactor array serialization into separate JAX and tensorstore logic

rdyro committed 321 days ago

0e1b3419

Makes Effort_02 the default value for memory_fitting_level.

Google-ML-Automation committed 321 days ago

0ae613ee

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

yashk2810 committed 321 days ago

80d1fbac

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

ZixuanJiang committed 321 days ago

d7c22eb1

Merge pull request #28063 from dfm:inline-array

Google-ML-Automation committed 322 days ago

c9fb5e57

Ensure outputs are tracers when inlining jit.

dfm committed 322 days ago

c6482ed6

Mark CliDebuggerTest as a thread-unsafe test.

hawkinsp committed 322 days 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 322 days ago

854b2c85

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

hawkinsp committed 322 days ago

934a0db4

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

hawkinsp committed 322 days ago

95156068

Allow the CPU collective implementation to be overridden to None.

hawkinsp committed 322 days ago

96865709

Reverts c2ba1790417ca206a4d88b25aef4d5ae510dd717

dfm committed 322 days ago

492cd3d9

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

hawkinsp committed 322 days ago

13686bd8

Newer Older