Repository: apache/tvm

Issue:
[Relay][Bug] nn.batch_norm not being simplified by SimplifyInference
cc @hypercubestart, who discovered this oddity

```python
r = """
#[version = "0.0.5"]
def @main(%data: Tensor[(1, 32, 32, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> (Tensor[(1, 32, 32, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) {
    nn.batch_norm(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var, axis=3) /* ty=(Tensor[(1, 32, 32, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) */
}
"""
mod = tvm.parser.parse(r)
mod = relay.transform.SimplifyInference()(mod)
print(mod.astext())
```
the result is not what we expected, though, as the batch norm doesn’t get simplified:
```
#[version = "0.0.5"]
def @main(%data: Tensor[(1, 32, 32, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> (Tensor[(1, 32, 32, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) {
  nn.batch_norm(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var, axis=3) /* ty=(Tensor[(1, 32, 32, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) */
}
```
is that to be expected? are we doing something dumb here?


Discussion:
Nevermind, I think this is to be expected. It's because we weren't indexing into the tuple result/we hadn't changed the return type of the function.

```python
import tvm
from tvm import relay

r = """
#[version = "0.0.5"]
def @main(%data: Tensor[(1, 32, 32, 16), float32], %bn_gamma: Tensor[(16), float32], %bn_beta: Tensor[(16), float32], %bn_mean: Tensor[(16), float32], %bn_var: Tensor[(16), float32]) -> Tensor[(1, 32, 32, 16), float32] {
    %0 = nn.batch_norm(%data, %bn_gamma, %bn_beta, %bn_mean, %bn_var, axis=3) /* ty=(Tensor[(1, 32, 32, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) */;
    %0.0
}
"""
mod = tvm.parser.parse(r)
mod = relay.transform.SimplifyInference()(mod)
print(mod.astext())
```
The root cause is that the simplify pass detects the `batch_norm -> get()` pattern, so it would ignore the batch_norm without get().
Makes sense. We hit a weird corner case here, glad we figured it out. Maybe a documentation change is warranted, but it's also a pretty niche corner case, so I'm not sure it matters!
It doesn't hurt to improve the document for this case as we met this issue as well several months ago lol
Please let me know if you would file a PR; otherwise I will put it in my TODO list but I'm not sure when I will do so..

Resolving patch:
diff --git a/python/tvm/relay/transform/transform.py b/python/tvm/relay/transform/transform.py
--- a/python/tvm/relay/transform/transform.py
+++ b/python/tvm/relay/transform/transform.py
@@ -173,10 +173,14 @@ def SimplifyInference():
     """Simplify the data-flow graph for inference phase. An simplified expression
     which is semantically equal to the input expression will be returned.
 
+    Note that batch norms will only be simplified if their result is indexed at
+    tuple index 0.
+
     Returns
     -------
     ret: tvm.transform.Pass
         The registered pass to perform operator simplification.
+
     """
     return _ffi_api.SimplifyInference()
 


Test patch: