[caffe2] Shape inference for UnPackRecords
Summary:
Since UnPackRecords is part of the graph, we need to add shape inference for it to make it work e2e with tvm_jit_op. Because the input is packed, shape inference is impossible without shape info of the packed tensors. Some context, the shape of the packed tensor is 1 X num_embeddings X embedding_size, with 1 being the batch_size. The shape of the corresponding output tensor is thus batch_size X num_embeddings X embedding_size after concatenating the packed tensors on the batch axis. Therefore two more gflags need to be added
- caffe2_predictor_num_embeddings
- caffe2_predictor_embedding_size
These gflags are then added to the UnPackRecordsOp in the predict_net as args to pass the info to c2_frontend so TVM can do its own shape inference.
Reviewed By: yinghai
Differential Revision: D21286983
fbshipit-source-id: e9a19cb6b564905282a771df2b9d211d5d37dd71
