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