I am doing something very similar where I need to replace BATCH_MATMUL operator with RESHAPE and MATMUL. were you successful in doing this? please share insights.