How to update batch normalization node's state variables?

The MPSCNNBatchNormalizationDataSource's optional methods updateGammaAndBetaWithCommandBuffer:batchNormalizationState: and updateMeanAndVarianceWithCommandBuffer:bathNormalizationState: should be used to update the state. To calculate the new state values (beta, gamma, mean and variance), MPSBatchNormalizationStatistics should be used, I suppose.


Naturally, MPSBatchNormalizationStatistics requires for the calculation of these new values "source images" (e.g. encodeBatchToCommandBuffer:sourceImage:batchNormalizationState:). The source images should be the source images of the batch normalization node.


But how do I access the input of the batch normalization node when updateGammaAndBetaWithCommandBuffer:batchNormalizationState: is called? I could not find any possibility to access these data. Any clues?