Systematic way of analyzing the program to make sure it is partitionable along the batch dimension

When we create micro-batches, we are assuming that the program can be split up along the specified batch dimension. We should have a way to analyze the program and tell if this is true or not (e.g. it would not be trivially true if we have something like BatchNorm, which has interactions across elements in the batch).