ROCm/jax

Pull Requests Commits

[Mosaic] Add support for dynamic filter values

naummo committed 327 days ago

7dda1731

Export `jax.smap` as a public API. `smap` is an ergonomic API to write single axis shard_maps and make nesting of shard maps delightful to express :)

yashk2810 committed 327 days ago

da1088b0

Update XLA dependency to use revision

Google-ML-Automation committed 327 days ago

90e18d5f

Allow async remote copy to take device id as a dict of axis names and indices.

IvyZX committed 327 days ago

c7a3b685

Expose `out_info` on Compiled just as it exists on `Lowered` and `Traced`.

yashk2810 committed 327 days ago

00dbcee8

Merge pull request #30078 from mattjj:smap

Google-ML-Automation committed 327 days ago

ee9e148a

Alter internal JAX calls to auto_axes to apply over only explicit axes. This avoids cases where auto_axes is used under a shard_map which leads to an error.

oliverdutton-iso committed 327 days ago

5ad539c3

[shard-map] fix partial manual eager mode

mattjj committed 327 days ago

3a8ba569

Allow `P(('a', 'b'))` in aval.sharding where mesh is `('a', 2, Explicit), ('b', 2, Auto)`.

yashk2810 committed 327 days ago

35009ab7

[Mosaic] Support all reshapes where minor dim % 128 == 0 for 16-bit and 8-bit types. Extend various tests for these cases.

jaswanth- committed 327 days ago

96797a9f

Move xla_workspace4() and xla_workspace3() calls above Python init.

Google-ML-Automation committed 327 days ago

b1341b11

Merge pull request #30082 from jakevdp:dep-xla-canonicalize

Google-ML-Automation committed 327 days ago

34b0e3a5

Add support for logging vectors on SC.

sashabu committed 327 days ago

61250a3c

Deprecate jax.interpreters.xla.canonicalize_dtype.

jakevdp committed 327 days ago

d551909c

Merge pull request #30079 from jakevdp:dep-compile-options

Google-ML-Automation committed 327 days ago

1a8adcb2

[mosaic] Clarified that tpu.device_id return a logical ID

superbobry committed 327 days ago

75d6129e

[pallas:mosaic] Enabled pytype for Mosaic TPU lowering

superbobry committed 327 days ago

339b58e4

[Pallas:MGPU] Add support for reductions that might require cross-warp communication

apaszke committed 327 days ago

5300b4de

[Mosaic GPU] Insert nvvm.minctasm = 1 to make sure ptxas recognizes setmaxnreg

apaszke committed 327 days ago

35ce3bc1

Finalize deprecation of jax.core.typecheck

Jake VanderPlas committed 327 days ago

ea508e35

Add `use_shardy_partitioner` in `tf.tpu.XlaOptions`, `TPUReplicateMetadata`. The default value is false.

ZixuanJiang committed 327 days ago

e9677067

[Mosaic GPU] Improve the error message when FragmentedArray slice is not tile aligned

apaszke committed 327 days ago

aa9fdd3c

[Pallas:MGPU] Create a programatic way to derive reduced layouts from tiled layouts in Pallas

apaszke committed 327 days ago

c66d3693

[mosaic_gpu] Enable pytype for Mosaic GPU .py files

superbobry committed 327 days ago

0cdbd6a7

[Mosaic GPU] Shorten the name of `equations.Expression` variants.

bchetioui committed 327 days ago

01d997b4

Update XLA dependency to use revision

Google-ML-Automation committed 328 days ago

e8f4457a

Merge pull request #29570 from carlosgmartin:gitignore_build_editable

Google-ML-Automation committed 328 days ago

fdac4e68

Deprecate jax.lib.xla_bridge.get_compile_options

jakevdp committed 328 days ago

793c6102

Merge pull request #30077 from jakevdp:cleanup-unused

Google-ML-Automation committed 328 days ago

12847df7

Simplify `as_manual_mesh` now that shard_map only tracks manual axes and not `auto` as it did previously.

yashk2810 committed 328 days ago

ac9c4c42

Older