| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| """Util functions/classes for jax primitive test harnesses.""" | |
| import contextlib | |
| import functools | |
| from typing import Optional | |
| import warnings | |
| import zlib | |
| from absl import logging | |
| from absl.testing import absltest | |
| from absl.testing import parameterized | |
| import jax | |
| from jax.experimental import jax2tf | |
| from jax.experimental.jax2tf.tests import jax2tf_limitations | |
| from jax.experimental.jax2tf.tests import primitive_harness | |
| import ml_dtypes | |
| import numpy as np | |
| import numpy.random as npr | |
| import tensorflow as tf | |
| _SUPPORTED_DTYPES = [np.float32] | |
| def _harness_matches(harness, group_name, dtype, params): | |
| if harness.group_name != group_name: | |
| return False | |
| if dtype is not None and harness.dtype != dtype: | |
| return False | |
| for key, value in params.items(): | |
| if harness.params.get(key, None) != value: | |
| return False | |
| return True | |
| _CRASH_LIST_PARMS = [] | |
| _DEFAULT_TOLERANCE = { | |
| jax.dtypes.float0: 0, | |
| np.dtype(np.bool_): 0, | |
| np.dtype(ml_dtypes.int4): 0, | |
| np.dtype(np.int8): 0, | |
| np.dtype(np.int16): 0, | |
| np.dtype(np.int32): 0, | |
| np.dtype(np.int64): 0, | |
| np.dtype(ml_dtypes.uint4): 0, | |
| np.dtype(np.uint8): 0, | |
| np.dtype(np.uint16): 0, | |
| np.dtype(np.uint32): 0, | |
| np.dtype(np.uint64): 0, | |
| np.dtype(ml_dtypes.float8_e4m3b11fnuz): 1e-1, | |
| np.dtype(ml_dtypes.float8_e4m3fn): 1e-1, | |
| np.dtype(ml_dtypes.float8_e5m2): 1e-1, | |
| np.dtype(ml_dtypes.bfloat16): 1e-2, | |
| np.dtype(np.float16): 1e-3, | |
| np.dtype(np.float32): 1e-6, | |
| np.dtype(np.float64): 1e-15, | |
| np.dtype(np.complex64): 1e-6, | |
| np.dtype(np.complex128): 1e-15, | |
| } | |
| def _dtype(x): | |
| if hasattr(x, "dtype"): | |
| return x.dtype | |
| elif type(x) in jax.dtypes.python_scalar_dtypes: | |
| return np.dtype(jax.dtypes.python_scalar_dtypes[type(x)]) | |
| else: | |
| return np.asarray(x).dtype | |
| def tolerance(dtype, tol=None): | |
| tol = {} if tol is None else tol | |
| if not isinstance(tol, dict): | |
| return tol | |
| tol = {np.dtype(key): value for key, value in tol.items()} | |
| dtype = jax.dtypes.canonicalize_dtype(np.dtype(dtype)) | |
| return tol.get(dtype, _DEFAULT_TOLERANCE[dtype]) | |
| def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=""): | |
| """Checks if two numpy arrays are all close given tolerances. | |
| Args: | |
| a: The array to check. | |
| b: The expected array. | |
| atol: Absolute tolerance. | |
| rtol: Relative tolerance. | |
| err_msg: The error message to print in case of failure. | |
| """ | |
| if a.dtype == b.dtype == jax.dtypes.float0: | |
| np.testing.assert_array_equal(a, b, err_msg=err_msg) | |
| return | |
| custom_dtypes = [ | |
| ml_dtypes.float8_e4m3b11fnuz, | |
| ml_dtypes.float8_e4m3fn, | |
| ml_dtypes.float8_e5m2, | |
| ml_dtypes.bfloat16, | |
| ] | |
| a = a.astype(np.float32) if a.dtype in custom_dtypes else a | |
| b = b.astype(np.float32) if b.dtype in custom_dtypes else b | |
| kw = {} | |
| if atol: | |
| kw["atol"] = atol | |
| if rtol: | |
| kw["rtol"] = rtol | |
| with np.errstate(invalid="ignore"): | |
| # TODO(phawkins): surprisingly, assert_allclose sometimes reports invalid | |
| # value errors. It should not do that. | |
| np.testing.assert_allclose(a, b, **kw, err_msg=err_msg) | |
| @contextlib.contextmanager | |
| def ignore_warning(**kw): | |
| with warnings.catch_warnings(): | |
| warnings.filterwarnings("ignore", **kw) | |
| yield | |
| def _has_only_supported_dtypes(harness): | |
| return True | |
| if harness.dtype not in _SUPPORTED_DTYPES: | |
| return False | |
| for key, value in harness.params.items(): | |
| if "dtype" in key and value not in _SUPPORTED_DTYPES: | |
| return False | |
| return True | |
| def is_sequence(x): | |
| try: | |
| iter(x) | |
| except TypeError: | |
| return False | |
| else: | |
| return True | |
| def primitives_parameterized(harnesses, *, one_containing: Optional[str] = None): | |
| """Decorator for tests. This is used to filter the tests. | |
| Args: | |
| harnesses: List of Harness objects to be filtered. | |
| one_containing: If set, only creates one test case for the provided name. | |
| Returns: | |
| A parameterized version of the test function with filtered set of harnesses. | |
| """ | |
| def _filter_harness(harness): | |
| # TODO(b/295369536) Put a limitations system in place so what's not covered | |
| # is explicit. | |
| if not harness.params.get("enable_xla", True): | |
| return False | |
| if one_containing is not None and one_containing not in harness.fullname: | |
| return False | |
| if not _has_only_supported_dtypes(harness): | |
| return False | |
| for crash_item in _CRASH_LIST_PARMS: | |
| if _harness_matches( | |
| harness, | |
| crash_item["group_name"], | |
| crash_item["dtype"], | |
| crash_item["params"], | |
| ): | |
| return False | |
| return True | |
| harnesses = filter(_filter_harness, harnesses) | |
| return primitive_harness.parameterized(harnesses, include_jax_unimpl=False) | |
| class JaxRunningTestCase(parameterized.TestCase): | |
| """A test case for JAX conversions.""" | |
| def setUp(self): | |
| super().setUp() | |
| # We use the adler32 hash for two reasons. | |
| # a) it is deterministic run to run, unlike hash() which is randomized. | |
| # b) it returns values in int32 range, which RandomState requires. | |
| self._rng = npr.RandomState(zlib.adler32(self._testMethodName.encode())) | |
| def rng(self): | |
| return self._rng | |
| def assert_all_close( | |
| self, | |
| x, | |
| y, | |
| *, | |
| check_dtypes=True, | |
| atol=None, | |
| rtol=None, | |
| canonicalize_dtypes=True, | |
| err_msg="", | |
| ): | |
| """Assert that x and y, either arrays or nested tuples/lists, are close.""" | |
| if isinstance(x, dict): | |
| self.assertIsInstance(y, dict) | |
| self.assertEqual(set(x.keys()), set(y.keys())) | |
| for k in x.keys(): | |
| self.assert_all_close( | |
| x[k], | |
| y[k], | |
| check_dtypes=check_dtypes, | |
| atol=atol, | |
| rtol=rtol, | |
| canonicalize_dtypes=canonicalize_dtypes, | |
| err_msg=err_msg, | |
| ) | |
| elif is_sequence(x) and not hasattr(x, "__array__"): | |
| self.assertTrue(is_sequence(y) and not hasattr(y, "__array__")) | |
| self.assertEqual(len(x), len(y)) | |
| for x_elt, y_elt in zip(x, y): | |
| self.assert_all_close( | |
| x_elt, | |
| y_elt, | |
| check_dtypes=check_dtypes, | |
| atol=atol, | |
| rtol=rtol, | |
| canonicalize_dtypes=canonicalize_dtypes, | |
| err_msg=err_msg, | |
| ) | |
| elif hasattr(x, "__array__") or np.isscalar(x): | |
| self.assertTrue(hasattr(y, "__array__") or np.isscalar(y)) | |
| if check_dtypes: | |
| self.assert_dtypes_match(x, y, canonicalize_dtypes=canonicalize_dtypes) | |
| x = np.asarray(x) | |
| y = np.asarray(y) | |
| self.assert_arrays_all_close( | |
| x, y, check_dtypes=False, atol=atol, rtol=rtol, err_msg=err_msg | |
| ) | |
| elif x == y: | |
| return | |
| else: | |
| raise TypeError((type(x), type(y))) | |
| def assert_arrays_all_close( | |
| self, x, y, *, check_dtypes=True, atol=None, rtol=None, err_msg="" | |
| ): | |
| """Assert that x and y are close (up to numerical tolerances).""" | |
| self.assertEqual(x.shape, y.shape) | |
| atol = max(tolerance(_dtype(x), atol), tolerance(_dtype(y), atol)) | |
| rtol = max(tolerance(_dtype(x), rtol), tolerance(_dtype(y), rtol)) | |
| _assert_numpy_allclose(x, y, atol=atol, rtol=rtol, err_msg=err_msg) | |
| if check_dtypes: | |
| self.assert_dtypes_match(x, y) | |
| def assert_dtypes_match(self, x, y, *, canonicalize_dtypes=True): | |
| if not jax.config.x64_enabled and canonicalize_dtypes: | |
| self.assertEqual( | |
| jax.dtypes.canonicalize_dtype(_dtype(x), allow_opaque_dtype=True), | |
| jax.dtypes.canonicalize_dtype(_dtype(y), allow_opaque_dtype=True), | |
| ) | |
| else: | |
| self.assertEqual(_dtype(x), _dtype(y)) | |
| class PrimitivesTest(JaxRunningTestCase): | |
| @primitives_parameterized( | |
| primitive_harness.all_harnesses, | |
| ) | |
| @ignore_warning( | |
| category=UserWarning, message="Using reduced precision for gradient.*" | |
| ) | |
| def test_prim(self, harness: primitive_harness.Harness): | |
| def _filter_limitation(limitation): | |
| return limitation.filter(device="cpu", dtype=harness.dtype, mode="compiled") | |
| limitations = tuple( | |
| filter( | |
| _filter_limitation, | |
| jax2tf_limitations.Jax2TfLimitation.limitations_for_harness(harness), | |
| ) | |
| ) | |
| devices = [] | |
| possible_backends = ["iree_cpu"] | |
| for backend in possible_backends: | |
| try: | |
| devices.extend(jax.devices(backend)) | |
| except RuntimeError: | |
| continue | |
| func_jax = harness.dyn_fun | |
| args = harness.dyn_args_maker(self.rng()) | |
| baseline_platform: Optional[str] = None | |
| baseline_results = None | |
| for d in devices: | |
| if baseline_platform is None or baseline_platform != d.platform: | |
| device_args = jax.tree_util.tree_map( | |
| lambda x: jax.device_put(x, d), args | |
| ) | |
| logging.info("Running harness on %s", d) | |
| with jax.jax2tf_associative_scan_reductions(True): | |
| res = func_jax(*device_args) | |
| if baseline_platform is None: | |
| baseline_platform = d.platform | |
| baseline_results = res | |
| else: | |
| logging.info( | |
| "Comparing results for %s and %s", baseline_platform, d.platform | |
| ) | |
| self.assert_all_close(baseline_results, res) | |
| if __name__ == "__main__": | |
| absltest.main() |