What to do with passes

I have a small collection of relay passes that I found useful and that might be interesting for more general use:

  • Merge consecutive transpose ops (useful for converted models when the other source framework had different conventions, I’m using PyTorch),
  • merge equal (shape) constants to enable CSE (also for converted models, but probably also our own),
  • specialize _like ops for static shapes to _to (and similar) ops (broadcast, collapse_sum, zeros, ones, reshape) - useful after taking gradients because can reduce the complexity of the graph immensely. Removing broadcast/collapse_sum when not needed,
  • remove *1 and +0 (with tensor ones and zeros, these appear in gradients),
  • spell out training LayerNorm (batch norm etc. could be done, too).

I implemented these (mostly) using the PatternMatching in Python. Moving them to TVM likely involves going to C++.

Are these useful and if so, would we want to automatically include them in optimization (and how)?

Best regards

Thomas

1 Like