jax

e750d7ea - Add option for debug print to be called on partitioned arguments rather than all-gather and print the entire logical arguments.