[GitHub] [incubator-tvm] gussmith23 edited a comment on issue #6852: [Relay][Bug] nn.batch_norm not being simplified by SimplifyInference

[GitBox <gi...@apache.org>]

gussmith23 edited a comment on issue #6852:
URL: https://github.com/apache/incubator-tvm/issues/6852#issuecomment-722054681


   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())
   ```


----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org