mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 20:06:05 +00:00
5001a21bad
The primitive_harness.py defines a set of about 7000 test harnesses, each with a JAX callable and a recipe for generating the arguments for the callable. Note that the test harness does not define any expected behavior. The test harnesses can be used in several kinds of tests. Initially these harnesses were designed to test the completeness of the jax2tf lowering: for each test harness we convert it to TF and then we test that the result of invoking it is the same as for JAX native. Since then we have found other uses of test harnesses. * E.g., shape_poly_test.py tests that we can apply `jax.vmap` to each test harness and that we get a JAX callable that can be traced shape polymorphically, using a dimension variable for the batch dimension. * E.g., multi_platform_lowering_test.py tests that we can generate multi-platform lowering for each test harnesse. * E.g., the TFLite team is using the test harnesses to check the completeness of the TFLite lowering. Since the test harnesses are useful for non-jax2tf uses we hereby moved them to jax._src.internal_test_util.test_harnesses. (We also renamed the module from primitive_harness to test_harnesses.) This change is necessary to move some tests out of jax2tf: multi_platform_lowering_test.py, shape_poly_test.py. PiperOrigin-RevId: 581016785
108 lines
3.1 KiB
TOML
108 lines
3.1 KiB
TOML
[build-system]
|
|
requires = ["setuptools", "wheel"]
|
|
build-backend = "setuptools.build_meta"
|
|
|
|
[tool.mypy]
|
|
show_error_codes = true
|
|
disable_error_code = "attr-defined, name-defined, annotation-unchecked"
|
|
no_implicit_optional = true
|
|
|
|
[[tool.mypy.overrides]]
|
|
module = [
|
|
"absl.*",
|
|
"colorama.*",
|
|
"importlib_metadata.*",
|
|
"IPython.*",
|
|
"numpy.*",
|
|
"opt_einsum.*",
|
|
"scipy.*",
|
|
"libtpu.*",
|
|
"jaxlib.mlir.*",
|
|
"iree.*",
|
|
"rich.*",
|
|
"optax.*",
|
|
"flax.*",
|
|
"tensorflow.*",
|
|
"tensorflowjs.*",
|
|
"tensorflow.io.*",
|
|
"tensorstore.*",
|
|
"web_pdb.*",
|
|
"etils.*",
|
|
"google.colab.*",
|
|
"pygments.*",
|
|
"jraph.*",
|
|
"matplotlib.*",
|
|
"tensorboard_plugin_profile.convert.*",
|
|
"jaxlib.*",
|
|
"pytest.*",
|
|
"zstandard.*",
|
|
"jax.experimental.jax2tf.tests.flax_models",
|
|
"jax.experimental.jax2tf.tests.back_compat_testdata",
|
|
"setuptools.*",
|
|
]
|
|
ignore_missing_imports = true
|
|
|
|
[[tool.mypy.overrides]]
|
|
module = [
|
|
"jax.interpreters.autospmd",
|
|
"jax.lax.lax_parallel",
|
|
"jax._src.internal_test_util.test_harnesses",
|
|
]
|
|
ignore_errors = true
|
|
|
|
[tool.pytest.ini_options]
|
|
markers = [
|
|
"multiaccelerator: indicates that a test can make use of and possibly requires multiple accelerators",
|
|
"SlurmMultiNodeGpuTest: mark a test for Slurm multinode GPU nightly CI"
|
|
]
|
|
filterwarnings = [
|
|
"error",
|
|
"ignore:No GPU/TPU found, falling back to CPU.:UserWarning",
|
|
"ignore:xmap is an experimental feature and probably has bugs!",
|
|
"ignore:the imp module is deprecated in favour of importlib.*:DeprecationWarning",
|
|
"ignore:can't resolve package from __spec__ or __package__:ImportWarning",
|
|
"ignore:Using or importing the ABCs.*:DeprecationWarning",
|
|
"ignore:numpy.ufunc size changed",
|
|
"ignore:.*experimental feature",
|
|
"ignore:The distutils.* is deprecated.*:DeprecationWarning",
|
|
"default:Error reading persistent compilation cache entry for 'jit__lambda_'",
|
|
"default:Error writing persistent compilation cache entry for 'jit__lambda_'",
|
|
"ignore:backend and device argument on jit is deprecated.*:DeprecationWarning",
|
|
# TODO(skyewm): remove when jaxlib >= 0.4.12 is released (needs
|
|
# https://github.com/openxla/xla/commit/fb9dc3db0999bf14c78d95cb7c3aa6815221ddc7)
|
|
"ignore:ml_dtypes.float8_e4m3b11 is deprecated.",
|
|
"ignore:JAX_USE_PJRT_C_API_ON_TPU=false will no longer be supported.*:UserWarning",
|
|
"ignore:np.find_common_type is deprecated.*:DeprecationWarning",
|
|
"ignore:jax.numpy.in1d is deprecated.*:DeprecationWarning",
|
|
"ignore:jax.numpy.trapz is deprecated.*:DeprecationWarning",
|
|
]
|
|
doctest_optionflags = [
|
|
"NUMBER",
|
|
"NORMALIZE_WHITESPACE"
|
|
]
|
|
addopts = "--doctest-glob='*.rst'"
|
|
|
|
[tool.pylint.master]
|
|
extension-pkg-whitelist = "numpy"
|
|
|
|
[tool.pylint."messages control"]
|
|
disable = [
|
|
"missing-docstring",
|
|
"too-many-locals",
|
|
"invalid-name",
|
|
"redefined-outer-name",
|
|
"redefined-builtin",
|
|
"protected-name",
|
|
"no-else-return",
|
|
"fixme",
|
|
"protected-access",
|
|
"too-many-arguments",
|
|
"blacklisted-name",
|
|
"too-few-public-methods",
|
|
"unnecessary-lambda"
|
|
]
|
|
enable = "c-extension-no-member"
|
|
|
|
[tool.pylint.format]
|
|
indent-string=" "
|