ImportError: cannot import name 'TensorCoreTiledLayoutType'

[andrewor14]

Hi all, I'm still hitting this issue locally even after #1886. It's because my local torchao version is "0.6.0+gitbd09e9d2", so it still fails. My local copy is on the latest torchao main but we haven't bumped the version yet, so it needs the new import but the version guard thinks otherwise.

I propose that we delete the import guard and simply do try catch around the imports instead, e.g.

try:
    # torchao 0.7+
    from torchao.dtypes import TensorCoreTiledLayout
catch ImportError:
    # torchao 0.6 and before
    from torchao.dtypes import TensorCoreTiledLayoutType as TensorCoreTiledLayout

This seems more robust to me, and I already have something similar for migrating QAT out of prototype in this PR: #1883

Thoughts? @jainapurva @joecummings @ebsmothers

[ebsmothers]

Yeah this is something that's come up before too.. the current version parsing is not really able to pick up the right version if you're running on a source install. I do like the ability to designate a version guard in a centralized place as I think it'll be more likely that we clean things up in a timely fashion (the code comment you've added is a decent approximation to this though). We could update the version guard to check for source installs, but that is a bit messier and the try/except is nice and simple. So yeah I'm good with the approach you're proposing here

[joecummings]

@andrewor14 Are you planning on submitting this as a PR? We have a code freeze tomorrow and would love to get this in

[andrewor14]

Yeah sure I'll submit a PR later today