I am not sure, but if you could open a github issue with a simple toy script that exercises the bmm pattern you care about, and also call out what specifically you’re wanting the compiler to do, we can check with different teams and see if something we already have does this today (e.g. torchdynamo) or can have support added.