Tensorflow layout tracking

Anyone knows how to track layout or axis settings for each node when building?

Traceback (most recent call last):
File “x.py”, line 142, in
graph, lib, params = nnvm.compiler.build(sym, shape=shape_dict, target=target, target_host=target_host, dtype=dtype_dict, params=params)
File “/homed/lidw/tvm/nnvm/python/nnvm/compiler/build_module.py”, line 270, in build
ishape, _ = graph_util.infer_shape(graph, **shape)
File “/homed/lidw/tvm/nnvm/python/nnvm/compiler/graph_util.py”, line 31, in infer_shape
graph = graph.apply(“InferShape”)
File “/homed/lidw/tvm/nnvm/python/nnvm/graph.py”, line 234, in apply
check_call(_LIB.NNGraphApplyPasses(self.handle, npass, cpass, ctypes.byref(ghandle)))
File “/homed/lidw/tvm/nnvm/python/nnvm/_base.py”, line 75, in check_call
raise NNVMError(py_str(_LIB.NNGetLastError()))
nnvm._base.NNVMError: Error in operator model_with_buckets/embedding_attention_decoder_6/attention_decoder/attention_decoder/MatMul: [21:51:24] /homed/lidw/tvm-zx/nnvm/src/top/nn/nn.cc:58: Operator dense(use_bias=False, units=128, name=model_with_buckets/embedding_attention_decoder_6/attention_decoder/attention_decoder/MatMul) expects weight’s shape to be [128,128], but got [128,384].