jax.experimental.multihost_utils.assert_equal

Contents

jax.experimental.multihost_utils.assert_equal#

jax.experimental.multihost_utils.assert_equal(in_tree, fail_message='')[source]#

Verifies that all the hosts have the same tree of values.

Parameters:

fail_message (str) –