How can I create a deeplearning4j Lambda Layer taking the mean over one dimension for a keras model inported from python? - deeplearning4j

Question

I try to load a Keras Model trained in Python into Java using deeplearning4j. The Model contains a Lambda Layer, that takes the average in one dimension. I found out, that I have to create the Lambda Layer as a Class, so that deeplearning4j knows what to expect. However, I need help with correctly setting the output dimensions of that Layer in Java.
I have tried to construct a InputType with the shape I want. However, it is an Abstract Class and the derived Classes have no Constructor allowing to set the Shape. I also didn't find any resize options.
The only thing I found, but that didn't help, was the inferInputType, but I could not get anything useful, as returnTypes is again an Array, while I need a single Type with 2 Dimensions...
long[] shape = inputType.getShape();
INDArray inputShapeA = Nd4j.zeros(shape[0]);
INDArray inputShapeB = Nd4j.zeros(shape[2]);
INDArray[] inputShape = {inputShapeA,inputShapeB};
InputTypes[] returnTypes = InputType.InputTypeFeedForward.inferInputTypes(inputShape)
b_features = 2 + 20 + 100 # (#temporal, #user, #doc)
##### Main part #####
inputs = Input(shape=(None, nb_features), name="Capture_Input")
##### Classifier #####
z = Lambda(lambda x: K.mean(x, axis=1), output_shape=lambda s: (s[0], s[2]), name="Lambda_Mean_Layer")(
inputs) # (None, nb_score)
outputs = Dense(1, activation='tanh', name="classification_Layer")(z)
##### Model #####
model = Model(input=[inputs], output=outputs, name="minimal_Import_Error_Model")
##### Compile #####
adam = Adam(lr=0.0001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
model.compile(optimizer=adam, loss='binary_crossentropy')
//training; could be done with random values, as the error is independent of the actual classification performance
model.summary()
model.save("model_name.h5")
Model Summary from Python, this is what I expected:
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
Capture_Input (InputLayer) (None, None, 122) 0
_________________________________________________________________
Lambda_Mean_Layer (Lambda) (None, 122) 0
_________________________________________________________________
classification_Layer (Dense) (None, 1) 123
=================================================================
Total params: 123
Trainable params: 123
Non-trainable params: 0
_________________________________________________________________
Java Code defining Lambda Layer (Error source), loading Model and getting random predictions.
import org.apache.log4j.BasicConfigurator;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.samediff.SameDiffLambdaLayer;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.modelimport.keras.KerasLayer;
import org.deeplearning4j.nn.modelimport.keras.KerasModelImport;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.autodiff.samediff.SDVariable;
import org.nd4j.autodiff.samediff.SameDiff;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.io.ClassPathResource;// load the model
class OurLambdaLayer extends SameDiffLambdaLayer {
// Python Keras version: Lambda(lambda x: K.mean(x, axis=1), output_shape=lambda s: (s[0], s[2]))(sub_h)
#Override
public SDVariable defineLayer(SameDiff sd, SDVariable x) {
//visited only once, right before the Error
System.out.println(x.toString());
for (int i = 0; i < x.getShape().length; i++) {
System.out.print("dim " + i + " :" + x.getShape()[i] + " ");
}
System.out.println("");
SDVariable returnvalue = x.mean(1);
System.out.println(returnvalue.toString());
return returnvalue;
}
#Override
public InputType getOutputType(int layerIndex, InputType inputType) {
// visited 2 times during Graph construction
return inputType;
}
}
class Main {
public static void main(String[] argv) {
BasicConfigurator.configure(); //some logging library initialization
try {
//prepare for lambda layer
KerasLayer.registerLambdaLayer("Lambda_Mean_Layer", new OurLambdaLayer());
// load the model
ClassPathResource cpr = new ClassPathResource("model_name.h5");
String simpleMlp = cpr.getFile().getPath();
ComputationGraph model = KerasModelImport.importKerasModelAndWeights(simpleMlp);
// make a random sample
int inputs = 122;
int[] inputA = {1, 1, inputs};
INDArray features = Nd4j.rand(inputA);
// get the prediction
INDArray[] prediction = model.output(features);
} catch (Exception e) {
e.printStackTrace();
}
}
}
I get the following Error, I think this is due to the Lambda Layer producing an Output which has different dimensions to what its getOutputType promised. I want to change getOutputType to correspond to the ouptut I create.
SDVariable(name="input",variableType=VARIABLE,dtype=FLOAT)
dim 0 :1 dim 1 :1 dim 2 :122
SDVariable(name="reduce_mean_1",variableType=ARRAY,dtype=FLOAT)
java.lang.IllegalArgumentException: Invalid input: expect NDArray with rank 3 (i.e., activations for RNN layer)
at org.deeplearning4j.nn.conf.preprocessor.RnnToFeedForwardPreProcessor.preProcess(RnnToFeedForwardPreProcessor.java:56)
at org.deeplearning4j.nn.graph.vertex.impl.LayerVertex.applyPreprocessorAndSetInput(LayerVertex.java:118)
at org.deeplearning4j.nn.graph.vertex.impl.LayerVertex.setInput(LayerVertex.java:170)
at org.deeplearning4j.nn.graph.ComputationGraph.outputOfLayersDetached(ComputationGraph.java:2388)
at org.deeplearning4j.nn.graph.ComputationGraph.output(ComputationGraph.java:1737)
at org.deeplearning4j.nn.graph.ComputationGraph.output(ComputationGraph.java:1693)
at org.deeplearning4j.nn.graph.ComputationGraph.output(ComputationGraph.java:1623)
at Main.main(main.java:74)