xla

cb5944ae - Fix the error where mark_step does not materalize tensors on SPMD:0 (#5281)

Fix the error where mark_step does not materalize tensors on SPMD:0 (#5281) * Fix the error where mark_step does not materalize tensors on SPMD:0 * typo * fix test_non_tensor_scalar