Unflatten traced module

.flake8

@@ -28,5 +28,5 @@ exclude =
     ./torch/include,
     ./torch/lib,
     ./venv,
-    ./pippy/fx,
+    ./pippy/unflatten.py,
     *.pyi

.github/workflows/pippy_tests.yaml

@@ -90,10 +90,16 @@ jobs:
           if [ -f requirements.txt ]; then pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html; fi
       - name: Install pippy
         run: "python setup.py install"
+      # IR tests
       - name: Test forward pipe generation
         run: python test/test_pipe.py
       - name: Test backward pipe generation
         run: python test/test_pipe_bwd.py
+      - name: Test unflatten
+        run: python test/test_unflatten.py
+      - name: Test Transformer
+        run: python test/test_transformer.py
+      # Runtime tests
       - name: Run forward-only integration test
         run: torchrun --nproc-per-node 4 test/test_fwd.py
       - name: Run forward-loss-backward integration test

check.sh

@@ -49,7 +49,7 @@ RETVAL=0
 
 if (( SKIP_FORMAT == 0 )); then
   echo; echo "Running format check ..."
-  ufmt diff pippy/*.py pippy/hf/*.py test/*.py
+  ufmt diff pippy/*.py test/*.py
   (( RETVAL |= $? ))
 fi
 

format.sh

@@ -7,7 +7,7 @@
 DEFAULT_TARGETS=()
 for f in $(git ls-files | grep '\.py$'); do
   case "$f" in
-    'pippy/fx/'*)
+    'pippy/unflatten.py')
       # ignore
       ;;