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 make some changes in SDY round trip #17923

Closed copybara-service[bot] closed 1 week ago

copybara-service[bot] commented 1 week ago

sdy make some changes in SDY round trip

  1. Make it a no-op for non SDY modules (i.e. don't add a meshes frontend attr)
  2. Fix bug in SDY round trip import that didn't take into account non inlined functions with result shardings.
  3. Add some missing tests for func return shardings.