def assertDtypesMatch(self, x, y, *, canonicalize_dtypes=True):
    'Compares dtypes across JAX and TF dtypes. Overrides super method.'

    def to_numpy_dtype(dt):
        return (dt if isinstance(dt, np.dtype) else dt.as_numpy_dtype)
    if ((not config.FLAGS.jax_enable_x64) and canonicalize_dtypes):
        self.assertEqual(dtypes.canonicalize_dtype(to_numpy_dtype(jtu._dtype(x))), dtypes.canonicalize_dtype(to_numpy_dtype(jtu._dtype(y))))
    else:
        self.assertEqual(to_numpy_dtype(jtu._dtype(x)), to_numpy_dtype(jtu._dtype(y)))
Compares dtypes across JAX and TF dtypes. Overrides super method.
def assertDtypesMatch(self, x, y, *, canonicalize_dtypes=True):
    

    def to_numpy_dtype(dt):
        return (dt if isinstance(dt, np.dtype) else dt.as_numpy_dtype)
    if ((not config.FLAGS.jax_enable_x64) and canonicalize_dtypes):
        self.assertEqual(dtypes.canonicalize_dtype(to_numpy_dtype(jtu._dtype(x))), dtypes.canonicalize_dtype(to_numpy_dtype(jtu._dtype(y))))
    else:
        self.assertEqual(to_numpy_dtype(jtu._dtype(x)), to_numpy_dtype(jtu._dtype(y)))