Correctly render the "layers" tree view

[kalekundert]

The best way to explain this PR is to show an example before and after:

Before:

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
GenotypeNetwork                               --                        --
├─Sequential: 1-1                             [2, 48, 32, 32]           --
│    └─Conv2d: 2-1                            [2, 48, 32, 32]           1,296
│    └─BatchNorm2d: 2-2                       [2, 48, 32, 32]           96
├─ModuleList: 1-2                             --                        --
│    └─Cell: 2-3                              [2, 32, 32, 32]           --
│    │    └─ReLUConvBN: 3-1                   [2, 32, 32, 32]           --
│    │    │    └─Sequential: 4-1              [2, 32, 32, 32]           --
│    │    │    │    └─ReLU: 5-1               [2, 48, 32, 32]           --
│    │    │    │    └─Conv2d: 5-2             [2, 32, 32, 32]           1,536
│    │    │    │    └─BatchNorm2d: 5-3        [2, 32, 32, 32]           64
│    │    └─ReLUConvBN: 3-2                   [2, 32, 32, 32]           --
│    │    │    └─Sequential: 4-2              [2, 32, 32, 32]           --
│    │    │    │    └─ReLU: 5-4               [2, 48, 32, 32]           --
│    │    │    │    └─Conv2d: 5-5             [2, 32, 32, 32]           1,536
│    │    │    │    └─BatchNorm2d: 5-6        [2, 32, 32, 32]           64
│    │    └─ModuleList: 3-3                   --                        --
│    │    │    └─IdentityModel: 4-3           [2, 32, 32, 32]           --
│    │    │    │    └─Identity: 5-7           [2, 32, 32, 32]           --
│    │    │    └─IdentityModel: 4-4           [2, 32, 32, 32]           --
│    │    │    │    └─Identity: 5-8           [2, 32, 32, 32]           --
===============================================================================================
Total params: 4,592
Trainable params: 4,592
Non-trainable params: 0
Total mult-adds (M): 8.95
===============================================================================================
Input size (MB): 0.02
Forward/backward pass size (MB): 3.67
Params size (MB): 0.02
Estimated Total Size (MB): 3.71
===============================================================================================

After:

===============================================================================================
Layer (type:depth-idx)                        Output Shape              Param #
===============================================================================================
GenotypeNetwork                               --                        --
├─Sequential: 1-1                             [2, 48, 32, 32]           --
│    ├─Conv2d: 2-1                            [2, 48, 32, 32]           1,296
│    └─BatchNorm2d: 2-2                       [2, 48, 32, 32]           96
└─ModuleList: 1-2                             --                        --
     └─Cell: 2-3                              [2, 32, 32, 32]           --
          ├─ReLUConvBN: 3-1                   [2, 32, 32, 32]           --
          │    └─Sequential: 4-1              [2, 32, 32, 32]           --
          │         ├─ReLU: 5-1               [2, 48, 32, 32]           --
          │         ├─Conv2d: 5-2             [2, 32, 32, 32]           1,536
          │         └─BatchNorm2d: 5-3        [2, 32, 32, 32]           64
          ├─ReLUConvBN: 3-2                   [2, 32, 32, 32]           --
          │    └─Sequential: 4-2              [2, 32, 32, 32]           --
          │         ├─ReLU: 5-4               [2, 48, 32, 32]           --
          │         ├─Conv2d: 5-5             [2, 32, 32, 32]           1,536
          │         └─BatchNorm2d: 5-6        [2, 32, 32, 32]           64
          └─ModuleList: 3-3                   --                        --
               ├─IdentityModel: 4-3           [2, 32, 32, 32]           --
               │    └─Identity: 5-7           [2, 32, 32, 32]           --
               └─IdentityModel: 4-4           [2, 32, 32, 32]           --
                    └─Identity: 5-8           [2, 32, 32, 32]           --
===============================================================================================
Total params: 4,592
Trainable params: 4,592
Non-trainable params: 0
Total mult-adds (M): 8.95
===============================================================================================
Input size (MB): 0.02
Forward/backward pass size (MB): 3.67
Params size (MB): 0.02
Estimated Total Size (MB): 3.71
===============================================================================================

Note that the extra vertical bars have been removed, and the ├ character is consistently used when the vertical bar should continue into the next line.

Minor comments:

• I used the walrus operator pretty extensively in formatting_test.py, so the tests won't pass for python 3.7. Technically python 3.7 is past its EOL, so I wrote the code the way I did in the hopes that you'll be willing to make 3.8 the minimum supported version. If not, let me know and I'll find another way to write those tests.

• I had to configure pytest to find files named *_test.py, otherwise the pre-commit check wouldn't pass. I'm worried that I'm doing something wrong, because I don't see how the pre-commit check could've worked for anyone before this, but presumably it did.

• While writing this code, I noticed that the "Param #" column for named parameters didn't respect the ascii_only setting; it would always use the unicode └ or ├ characters. Since I was changing a lot of the relevant code anyways, I fixed this bug while I was at it.

• I changed the ASCII-only format from:

  + Sequential
  |     + Conv2d
  |     + BatchNorm2d
  |     + ReLU
  |     + Conv2d
  |     + BatchNorm2d
  |     + ReLU
  

  to:

  |--Sequential
  |     |--Conv2d
  |     |--BatchNorm2d
  |     |--ReLU
  |     |--Conv2d
  |     |--BatchNorm2d
  |     '--ReLU
  

  My main motivation was to make the "Param #" column less confusing. With the original format, this column ends up just looking like a bunch of positive numbers, rather than a tree view.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (80d3e67) 97.50% compared to head (9fd3a13) 97.67%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #281      +/-   ##
==========================================
+ Coverage   97.50%   97.67%   +0.16%     
==========================================
  Files           6        6              
  Lines         641      687      +46     
==========================================
+ Hits          625      671      +46     
  Misses         16       16              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kalekundert]

I made some changes unrelated to my PR to get all of the tests (except the python 3.7 ones) to pass:

• I replaced # type: [import] with # type: [import-untyped], as apparently required by current versions of mypy.
• I added some type stub packages to requirements-dev.txt, so that pre-commit run -a completes without error.
• I skipped the FlanT5Small test for pytorch < 1.10. Previously pytorch < 1.8 was the cutoff, but it seems like a newer version of pytorch is now required.

[TylerYep]

Thanks for the PR! I would likely gate this behind a row setting at least to begin with before officially making this the default behavior.

[kalekundert]

For what it's worth, I don't think it makes sense to gate this behind a setting. Is the idea just to be gradual about making changes, or is there some way in which the current tree view is superior to the proposed one that I'm not appreciating?

[TylerYep]

More about being gradual with changes - some people may rely on the output matching exactly so there needs to be at least one version announcing the change before we switch. There may also be some bugs with this implementation that we can fix in between releases.

Apologies for taking so long to get to this PR, I haven't had the time recently to get this project up to date. Will get to a thorough review and merge soon (I need to fix some pre-commit stuff first).

[kalekundert]

Ok, that makes sense. Let me know if you want me to add the option, or if you'd rather do that yourself.