Skip to content

Commit

Permalink
[jax2tf] Add shard_map tests
Browse files Browse the repository at this point in the history 
Also fix tests to run on multiple devices in TF

PiperOrigin-RevId: 508691872
  • Loading branch information
@gnecula
gnecula authored and jax authors committed Feb 10, 2023
1 parent d09f3c2 commit be21404
Showing 1 changed file with 268 additions and 107 deletions.

0 comments on commit be21404

Please sign in to comment.