Export BatchNormalizationLayer as MXNet

I found the following workaround for NetChains with several Batchnorm layers by manually exporting aux states:

Mathematica

GetAux[layer_BatchNormalizationLayer] := {NetExtract[layer, 
      "MovingMean"], NetExtract[layer, "MovingVariance"]}
net = NetInitialize@NetChain@
     {BatchNormalizationLayer["Input" -> {2, 2, 2}], BatchNormalizationLayer["Input" -> {2, 2, 2}]}
aux = Flatten[GetAux /@ NetExtract[net, All], 1]
Export[FileNameJoin@{NotebookDirectory[],"batchnorm.json"}, net, "MXNet"]
Export[FileNameJoin@{NotebookDirectory[], "auxparam.txt"},
      ExportString[aux, "PythonExpression"]]

Then you can call the following lines of python code to import the aux states and execute the net:

import mxnet as mx
import numpy as np
sym=mx.symbol.load('batchnorm.json')
nd=mx.nd.load('batchnorm.params')

#get aux states from file
foo=np.loadtxt('auxparam.txt',dtype='str')
batchaux=eval(foo.item())
aux = {}
naux=sym.list_auxiliary_states()
for i in range(len(naux)):
    aux[naux[i]] = mx.nd.array(np.array(batchaux[i]))

#parse Input
inputND = mx.nd.array(np.array([[[[1,1],[2,2]]]]));
nd["Input"] = inputND

#run the net on input
e = sym.bind(mx.cpu(), nd, aux_states=aux);
out = e.forward()

This is fixed in Mathematica 12.0:

In[287]:= Export["~/batchnorm.json", 
foo = NetInitialize@BatchNormalizationLayer["Input" -> {1, 2, 2}], "MXNet"];
MXNetLink`NDArrayImport@"~/batchnorm.params"

Out[288]=
Association["Scaling" -> NumericArray[{1.}, "Real32"], 
 "Biases" -> NumericArray[{0.}, "Real32"], 
 "MovingMean" -> NumericArray[{0.}, "Real32"], 
 "MovingVariance" -> NumericArray[{1.}, "Real32"]]