Fix imports to get tests passing.

PiperOrigin-RevId: 274002692

oss_scripts/oss_pip_install.sh

@@ -3,22 +3,6 @@
 set -v  # print commands as they're executed
 set -e  # fail and exit on any command erroring
 
-: "${TF_VERSION:?}"
-
-if [[ "$TF_VERSION" == "tf-nightly"  ]]
-then
-  pip install tf-nightly;
-else
-  pip install -q "tensorflow==$TF_VERSION"
-fi
-
-# Make sure we have the latest version of numpy - avoid problems we were
-# seeing with Python 3
-pip install -q -U numpy
-
 # First ensure that the base dependencies are sufficient for a full import
 pip install -q -e .
 python -c "import mesh_tensorflow as mtf"
-
-# Then install the test dependencies
-pip install -q -e .[tests]

oss_scripts/oss_tests.sh

@@ -15,7 +15,7 @@ function set_status() {
 }
 
 # Run tests
-pytest --ignore=mesh_tensorflow/experimental
+python setup.py test
 set_status
 
 exit $STATUS

setup.cfg

@@ -0,0 +1,5 @@
+[aliases]
+test=pytest
+
+[tool:pytest]
+addopts=--ignore=mesh_tensorflow/experimental

setup.py

@@ -18,20 +18,19 @@
     },
     scripts=[],
     install_requires=[
+        'absl-py',
         'future',
         'gin-config',
+        'ortools',
         'six',
+        'tensorflow',
+        'tensorflow-datasets',
     ],
-    extras_require={
-        'auto_mtf': ['ortools>=7.0.6546'],
-        'tensorflow': ['tensorflow>=1.9.0'],
-        'tensorflow_gpu': ['tensorflow-gpu>=1.9.0'],
-        'tests': [
-            'absl-py',
-            'pytest',
-            'ortools>=7.0.6546',
-        ],
-    },
+    extras_require={},
+    tests_require=[
+        'pytest'
+    ],
+    setup_requires=['pytest-runner'],
     classifiers=[
         'Development Status :: 4 - Beta',
         'Intended Audience :: Developers',