File size: 2,887 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
"""
Runs tests that are appropriate for framework.
"""
import os
import sys
from subprocess import Popen
from pathlib import Path
__author__ = "Alex Rogozhnikov"
def run(cmd, **env):
# keeps printing output when testing
cmd = cmd.split(" ") if isinstance(cmd, str) else cmd
print("running:", cmd)
p = Popen(cmd, cwd=str(Path(__file__).parent), env={**os.environ, **env})
p.communicate()
return p.returncode
def main():
_executable, *args = sys.argv
frameworks = [x for x in args if x != "--pip-install"]
pip_install_is_set = "--pip-install" in args
framework_name2installation = {
"numpy": ["numpy"],
"torch": ["torch --index-url https://download.pytorch.org/whl/cpu"],
"jax": ["jax[cpu]", "flax"],
"tensorflow": ["tensorflow"],
"cupy": ["cupy"],
# switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927
# "paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"],
"paddle": ["paddlepaddle"],
"oneflow": ["oneflow==0.9.0"],
"pytensor": ["pytensor"],
}
usage = f"""
Usage: python -m einops.tests.run_tests <frameworks> [--pip-install]
Example: python -m einops.tests.run_tests numpy pytorch --pip-install
Available frameworks: {list(framework_name2installation)}
When --pip-install is set, auto-installs requirements with pip.
(make sure which pip points to right pip)
"""
if len(frameworks) == 0:
print(usage)
return
else:
synonyms = {
"tf": "tensorflow",
"pytorch": "torch",
"paddlepaddle": "paddle",
}
frameworks = [synonyms.get(f, f) for f in frameworks]
wrong_frameworks = [f for f in frameworks if f not in framework_name2installation]
if wrong_frameworks:
print(usage)
raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}")
if pip_install_is_set:
print("Install testing infra")
other_dependencies = ["pytest"]
assert 0 == run("pip install {} --progress-bar off -q".format(" ".join(other_dependencies)))
for framework in frameworks:
print(f"Installing {framework}")
pip_instructions = framework_name2installation[framework]
assert 0 == run("pip install {} --progress-bar off -q".format(" ".join(pip_instructions)))
# we need to inform testing script which frameworks to use
# this is done by setting an envvar EINOPS_TEST_BACKENDS
from einops.tests import unparse_backends
envvar_name, envvar_value = unparse_backends(backend_names=frameworks)
return_code = run(
"python -m pytest .",
**{envvar_name: envvar_value},
)
assert return_code == 0
if __name__ == "__main__":
main()
|