90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
"""
|
|
Runs tests that are appropriate for framework.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from subprocess import Popen
|
|
|
|
__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"],
|
|
}
|
|
if sys.platform == "darwin":
|
|
framework_name2installation["mlx"] = ["mlx"]
|
|
if sys.platform.startswith("linux"):
|
|
framework_name2installation["mlx"] = ["mlx[cpu]"]
|
|
|
|
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()
|