search

openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.64k stars 418 forks source link

#sdy Add support for inlined meshes in sdy round trip. #17970

Closed copybara-service[bot] closed 5 days ago

copybara-service[bot] commented 1 week ago

sdy Add support for inlined meshes in sdy round trip.

In addition, use sharding walker in -xla-sdy-round-trip-remove-size-one-axes, and fix issue where device ids were discarded from meshes that had size 1 axes.