'Get the dimension that corresponds to the batch size.
When data parallelism is used, the data will be automatically split and
concatenated along the batch-size dimension. Axis can be -1, which means
the whole array will be copied for each data-parallelism device.
Parameters
layout : str
layout string. For example, "NCHW".
Returns
int
An axis indicating the batch_size dimension.'

@staticmethod
def get_batch_axis(layout):

    if (layout is None):
        return 0
    return layout.find('N')
