Method assertNestedAllClose

axlearn/common/test_utils.py:244–262

(self, a, b, atol=1e-6, rtol=1e-3)

        return test_outputs, ref_outputs

    def assertNestedAllClose(self, a, b, atol=1e-6, rtol=1e-3):
        a_items = sorted(flatten_items(a), key=lambda x: x[0])
        b_items = sorted(flatten_items(b), key=lambda x: x[0])
        self.assertEqual([name for name, _ in a_items], [name for name, _ in b_items])
        for (a_name, a_value), (b_name, b_value) in zip(a_items, b_items):
            self.assertEqual(a_name, b_name)
            if isinstance(a_value, Tensor):
                a_value = as_local_tensor(a_value)
            if isinstance(b_value, Tensor):
                b_value = as_local_tensor(b_value)
            if isinstance(a_value, (np.ndarray, jnp.ndarray)) or isinstance(
                b_value, (np.ndarray, jnp.ndarray)
            ):
                a_value, b_value = as_tensor(a_value), as_tensor(b_value)
                self.assertEqual(a_value.dtype, b_value.dtype, msg=f"{a_name}")
                self.assertEqual(a_value.shape, b_value.shape, msg=f"{a_name}")
                assert_allclose(a_value, b_value, atol=atol, rtol=rtol, err_msg=f"{a_name}")
            else:
                self.assertAlmostEqual(a_value, b_value, msg=f"{a_name}")

    def assertNestedEqual(self, a, b):
        a_kv = flatten_items(a)