|
""" |
|
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): |
|
|
|
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"], |
|
|
|
|
|
"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))) |
|
|
|
|
|
|
|
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() |
|
|