def test_lstm_block_fused_op(
        self, compute_unit, tf_raw_lstm_op, is_flexible_input, batch_size, backend
    ):
        """
        Define a model with custom LSTM ops that uses tf.raw_ops.BlockLSTM / tf.raw_ops.BlockLSTMV2
        and verify that it converts to a fused lstm op.

        %x (shape: (Seq, Batch, idim) == (seq_len, batch, 4))
        %x1 = LSTM(h=10) (%input) # shape = (seq_len, batch, 10)
        %x2 = LSTM(h=20) (%x1) # shape = (seq_len, batch, 20)
        %x3 = slice()(%x2) # shape = (1, batch, 20), to get the final seq value
        %x4 = reshape((1, -1)) (%x3) # shape = (1, batch * 20)
        %x5 = Dense(h=3)(%x4) # shape = (1, 3)
        """

        class CustomLSTM(tf.keras.layers.Layer):
            def __init__(self, num_units, max_seq_length, batch_size):
                super(CustomLSTM, self).__init__()
                self.hidden_dim = num_units
                self.seq_length = max_seq_length
                self.batch_size = batch_size

            def build(self, input_shape):
                input_dim = input_shape[-1]
                self.w = self.add_weight(
                    shape=(input_dim + self.hidden_dim, 4 * self.hidden_dim),
                    initializer="random_normal",
                    trainable=True,
                )
                self.b = self.add_weight(shape=(4 * self.hidden_dim,), initializer="random_normal", trainable=True)
                self.init_h = tf.constant(np.zeros((self.batch_size, self.hidden_dim)).astype(np.float32))
                self.init_c = tf.constant(np.zeros((self.batch_size, self.hidden_dim)).astype(np.float32))

            def call(self, inputs):
                _, output_state, _, _, _, _, output = tf_raw_lstm_op(
                    seq_len_max=self.seq_length,
                    x=inputs,
                    cs_prev=self.init_c,
                    h_prev=self.init_h,
                    w=self.w,
                    wci=tf.constant(np.zeros((self.hidden_dim)).astype(np.float32)),
                    wcf=tf.constant(np.zeros((self.hidden_dim)).astype(np.float32)),
                    wco=tf.constant(np.zeros((self.hidden_dim)).astype(np.float32)),
                    b=self.b,
                )
                return output, output_state

        input_dim = 4
        seq_length = 5
        batch_size = batch_size
        x_shape = (seq_length, batch_size, input_dim)
        hidden_dim_1 = 10
        hidden_dim_2 = 20

        x = tf.keras.Input(batch_input_shape=x_shape)  # (seq_len, batch, 4)
        x1, output_states_1 = CustomLSTM(num_units=hidden_dim_1, max_seq_length=seq_length, batch_size=batch_size)(x)  # (seq_len, batch, 10), (seq_len, batch, 10)
        x2, output_states_2 = CustomLSTM(num_units=hidden_dim_2, max_seq_length=seq_length, batch_size=batch_size)(x1)  # (seq_len, batch, 20), (seq_len, batch 10)
        x3 = tf.slice(x2, begin=[4, 0, 0], size=[1, batch_size, 20])  # (1, batch, 20)
        x4 = tf.reshape(x3, shape=(1, -1))  # (1, batch * 20)
        x5 = tf.keras.layers.Dense(3)(x4)  # (1, 3)

        # Test that we can fuse the lstm op if we have an output that only extract the information from the last cell state
        x6 = tf.keras.layers.ReLU()(output_states_1[4, :, :])
        x7 = output_states_2[4:5, :, :]
        x8 = output_states_1[-1, :, :]
        x9 = tf.keras.layers.ReLU()(output_states_2[-1:, :, :])
        outputs = [x5, x8, x9] if is_flexible_input else [x5, x6, x7, x8, x9]

        keras_model = tf.keras.Model(inputs=x, outputs=outputs)

        inputs = None
        if is_flexible_input:
            inputs = [
                ct.TensorType(
                    shape=(ct.RangeDim(seq_length, 20), batch_size, input_dim)
                )
            ]

        res = TensorFlowBaseTest.run_compare_tf_keras(
            keras_model,
            [random_gen(x_shape, -1, 1)],
            compute_unit=compute_unit,
            backend=backend,
            inputs_for_conversion=inputs,
        )
        coreml_model = res[1]
        mil_prog = coreml_model._get_mil_internal()
        # assert that "lstm" ops are present in the mil program
        assert len(mil_prog.find_ops(op_type="lstm")) == 2

        Define a model with custom LSTM ops that uses tf.raw_ops.BlockLSTM / tf.raw_ops.BlockLSTMV2
        and verify that it converts to a fused lstm op.

        %x (shape: (Seq, Batch, idim) == (seq_len, batch, 4))
        %x1 = LSTM(h=10) (%input) # shape = (seq_len, batch, 10)
        %x2 = LSTM(h=20) (%x1) # shape = (seq_len, batch, 20)
        %x3 = slice()(%x2) # shape = (1, batch, 20), to get the final seq value
        %x4 = reshape((1, -1)) (%x3) # shape = (1, batch * 20)
        %x5 = Dense(h=3)(%x4) # shape = (1, 3)
        
coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops_tf_keras.py