BlockLSTM output values differ from Python

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary):
• TensorFlow version (use command below): 2.3.1
• Python version: 3.8
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory:

You can collect some of this information using our environment capture script
You can also obtain the TensorFlow version with
python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"

Describe the current behavior
The output of BlockLSTM differs from the one on Python.

Describe the expected behavior
I expected the output from BlockLSTM to have the same outputs values as Python

Code to reproduce the issue
I am using TF Java 0.2.0 and created a simple spike with hardcoded values in Java like this:

public static void main(String[] args) {
        EagerSession session = EagerSession.create();
        Ops tf = Ops.create(session);
        Scope scope = new Scope(session);

        float[][][] rawInputSequence = {{{0.1f, 0.2f}}, {{0.3f, 0.4f}}};
        Operand<TFloat32> inputSequence = tf.constant(rawInputSequence);

        int cellSize = 3;
        int[] cellShape = {1, cellSize};
        Operand<TInt32> cellDims = Constant.vectorOf(scope, cellShape);
        Operand<TInt64> seqLenMax = tf.array(2L);
        Operand<TFloat32> initialCellState = Zeros.create(scope, cellDims, TFloat32.DTYPE);
        Operand<TFloat32> initialHiddenState = Zeros.create(scope, cellDims, TFloat32.DTYPE);
        long[] biasShape = {cellSize * 4};
        Operand<TInt64> biasDim = Constant.vectorOf(scope, biasShape);
        Operand<TFloat32> bias = Zeros.create(scope, biasDim, TFloat32.DTYPE);

        FloatNdArray matrix = StdArrays.ndCopyOf(new float[][]{
            {1.6652163f, 1.366376f, 0.7786316f, 0.9834321f, 1.6551187f, -0.6363001f, -0.4229284f, 0.63195646f, 0.6605189f, -0.6906152f, 3.1515226f, 1.970373f},
            {1.9458166f, 0.9790728f, 0.7476161f, -1.6813406f, -0.75150734f, 0.13104685f, 0.004470979f, 0.009482844f, -1.1464607f, 0.5036645f, 1.3567412f, 0.71478313f},
            {0.5393334f, -0.6881541f, 1.5186735f, 1.3431606f, -0.61521095f, -2.1862414f, 1.2603592f, -0.33593372f, -0.48804748f, -0.34496853f, -0.8777565f, 0.9202126f},
            {1.3439888f, 0.32253885f, -0.7401764f, 0.10057431f, -1.3759913f, 0.08382488f, 0.56741005f, 2.207029f, -0.0066946335f, -0.8636334f, 1.9623716f, 0.14416508f},
            {-0.925145f, 0.2283957f, 0.79638815f, 0.2288384f, 0.7052175f, -0.18524477f, -2.308545f, 1.2240901f, 2.014674f, 0.6235778f, -0.15852839f, 0.17711076f}
        });
        Operand<TFloat32> weightMatrix = Constant.tensorOf(scope, matrix);

        FloatNdArray gatesMatrix = StdArrays.ndCopyOf(new float[]{1.6652163f, 1.366376f, 0.7786316f});
        Operand<TFloat32> weighGates = Constant.tensorOf(scope, gatesMatrix);

        BlockLSTM<TFloat32> blockLSTM = BlockLSTM.create(scope, seqLenMax, inputSequence, initialCellState,
                initialHiddenState, weightMatrix, weighGates, weighGates, weighGates, bias);

}

Whereas for Python I created the code below:

input_sequence = tf.constant([[[0.1, 0.2]], [[0.3, 0.4]]])
cell_size = 3
ini_cell_state = tf.zeros(shape=[1, cell_size])
ini_hidden_state = tf.zeros(shape=[1, cell_size])
bias = tf.zeros(shape=[cell_size * 4])
seq_len_max = tf.constant([2], dtype="int64")

weight_matrix = tf.constant([
    [1.6652163, 1.366376, 0.7786316, 0.9834321, 1.6551187, -0.6363001, -0.4229284, 0.63195646, 0.6605189, -0.6906152, 3.1515226, 1.970373],
    [1.9458166, 0.9790728, 0.7476161, -1.6813406, -0.75150734, 0.13104685, 0.004470979, 0.009482844, -1.1464607, 0.5036645, 1.3567412, 0.71478313],
    [0.5393334, -0.6881541, 1.5186735, 1.3431606, -0.61521095, -2.1862414, 1.2603592, -0.33593372, -0.48804748, -0.34496853, -0.8777565, 0.9202126],
    [1.3439888, 0.32253885, -0.7401764, 0.10057431, -1.3759913, 0.08382488, 0.56741005, 2.207029, -0.0066946335, -0.8636334, 1.9623716, 0.14416508],
    [-0.925145, 0.2283957, 0.79638815, 0.2288384, 0.7052175, -0.18524477, -2.308545, 1.2240901, 2.014674, 0.6235778, -0.15852839, 0.17711076]
])

weight_gates = tf.constant([1.6652163, 1.366376, 0.7786316])

block_lstm = tf.raw_ops.BlockLSTM(seq_len_max=seq_len_max, x=input_sequence, cs_prev=ini_cell_state,
                                  h_prev=ini_hidden_state, w=weight_matrix,
                                  wci=weight_gates, wcf=weight_gates,
                                  wco=weight_gates, b=bias)

Other info / logs
Output values for Java:

Input Gate: [0.6354535, 0.5823559, 0.5566029, 0.7944415, 0.6913817, 0.61125195]
Cell State: [-0.026291894, 0.037853386, -0.09006144, -0.016073283, 0.14818707, -0.25373507]
Forget State: [0.4407978, 0.50380254, 0.49064595, 0.4001113, 0.5333729, 0.4756381]
Output Gate: [0.50791717, 0.6425618, 0.58418906, 0.48629615, 0.825764, 0.70244145]
Cell Input: [-0.041375007, 0.06500043, -0.16180556, -0.0069905715, 0.18513232, -0.34502697]
Cell Output: [-0.026285835, 0.03783531, -0.08981872, -0.016071897, 0.1471118, -0.24842645]
Hidden Output: [-0.013351027, 0.024311526, -0.052471116, -0.007815702, 0.12147963, -0.17450504]

Output values for Python:

Input Gate: [[[0.6354535 0.5823559 0.5566029]]
 [[0.7784116 0.7010059 0.5999128]]]
Cell State: [[[-0.14840049  0.00885719 -0.02081871]]
 [[-0.45026863  0.16232333  0.0025551 ]]]
Forget State: [[[0.7228417  0.7436625  0.69778234]]
 [[0.6925149  0.7713927  0.67951703]]]
Output Gate: [[[0.50791717 0.6425618  0.58418906]]
 [[0.501899   0.8273453  0.69143474]]]
Cell Input: [[[-0.23353477  0.01520923 -0.03740317]]
 [[-0.44642073  0.2218112   0.02784033]]]
Cell Output: [[[-0.1473206   0.00885695 -0.02081571]]
 [[-0.4221198   0.1609125   0.0025551 ]]]
Hidden Output: [[[-0.07482666  0.00569114 -0.01216031]]
 [[-0.2118615   0.13313021  0.00176668]]]

Commencez par comparer l'appel Java à BlockLSTM.create avec tf.raw_ops.BlockLSTM de Python, en utilisant les entrées, les poids et les sorties codés en dur dans l'issue. Reproduisez la divergence dans les deux environnements et déterminez si le binding Java correspond à l'opération Python ; c'est terminé lorsque les valeurs de sortie correspondantes concordent.