def _try_replace_with_core_lstm(op: Operation) -> bool:
    """
    Inputs:

    op (Operation): op.op_type must be 'tf_lstm_block_cell' or `tf_lstm_block`

    Returns:

    True if op can be represented by mb.lstm op in SSA. False otherwise
    """
    def _check_unsupported_outputs(unsupported_outputs):
        for ov in unsupported_outputs:
            if len(ov.child_ops) > 0 or len(ov.consuming_blocks) > 0:
                return False
        return True

    # Check for unsupported configuration : When peephole is present
    if op.use_peephole.val:
        return False

    # Check if the tf lstm op can be replaced with coreml lstm op
    # We check the following two conditions
    # (1) The outputs must not be (i, f, o, ci, co), since there is no corresponding outputs with the LSTM in Core ML
    # (2) For the tf_lstm_block op, only the last time step of cell state can be used
    #     Here is an example of valid supported configuration:
    #          _, cell_states, _, _, _, _, _, _ = tf_lstm_block.outputs
    #          output = cell_states[-1, 1:2, :]
    #     And here is an example that coreml cannot handle currently:
    #          _, cell_states, _, _, _, _, _, _ = tf_lstm_block.outputs
    #          output = cell_states[:2, :, :]
    i, cs, f, o, ci, co, h = op.outputs
    unsupported_outputs = [i, f, o, ci, co]
    if not _check_unsupported_outputs(unsupported_outputs):
        return False

    if op.op_type == "tf_lstm_block":
        cs = _try_get_last_cell_state_in_tf_lstm_block(op)
        if cs is None:
            return False

    # op is compatible with lstm

    mb_peep = None
    if op.use_peephole.val:
        mb_peep = np.stack(
            [op.weight_peep_i.val, op.weight_peep_f.val, op.weight_peep_o.val]
        )

    # Set weights. The layout of the weight in TF1 is icfo (input, cell, forget, output gate).
    # Need to convert to ifoc for coreml
    tf_w = op.weight.val  # [input_dim+hidden_dim, 4*hidden_dim] in icfo layout
    tf_w_i, tf_w_c, tf_w_f, tf_w_o = np.split(tf_w, 4, axis=1)
    w = np.concatenate([tf_w_i, tf_w_f, tf_w_o, tf_w_c], axis=1)
    w = np.transpose(w, [1, 0])
    hidden_dim = w.shape[0] // 4
    input_dim = w.shape[1] - hidden_dim
    # Split input and hidden weights
    w_ih, w_hh = np.split(w, [input_dim], axis=1)

    # Bias is icfo. Convert to ssa LSTM's ifoc layout
    tf_b = op.bias.val
    tf_b_i, tf_b_c, tf_b_f, tf_b_o = np.split(tf_b, 4, axis=0)
    tf_b_f += op.forget_bias.val  # add forget bias to bias
    bias = np.concatenate([tf_b_i, tf_b_f, tf_b_o, tf_b_c], axis=0)

    cell_clip = None if op.cell_clip is None else op.cell_clip.val

    output_sequence = op.op_type == "tf_lstm_block"

    block = op.enclosing_block
    # x: [seq_len, batch, input_dim]
    if op.op_type == "tf_lstm_block_cell":
        x = mb.expand_dims(x=op.x, axes=[0], before_op=op)
    elif op.op_type == "tf_lstm_block":
        x = op.x
    else:
        raise ValueError("tf lstm op {} not supported. Only {} supported".format(op.op_type, SUPPORTED_TF_LSTM_OPS))

    new_h_all, new_h, new_cs = mb.lstm(
        x=x,
        initial_c=op.c_prev,
        initial_h=op.h_prev,
        weight_ih=w_ih,
        weight_hh=w_hh,
        bias=bias,
        recurrent_activation="sigmoid",
        cell_activation="tanh",
        activation="tanh",
        peephole=mb_peep,
        clip=cell_clip,
        output_sequence=output_sequence,
        name=op.name,
        before_op=op,
    )

    ops_to_remove = [op]
    block.replace_uses_of_var_after_op(anchor_op=op, old_var=cs, new_var=new_cs)
    if op.op_type == "tf_lstm_block_cell":
        block.replace_uses_of_var_after_op(anchor_op=op, old_var=h, new_var=new_h)
    elif op.op_type == "tf_lstm_block":
        block.replace_uses_of_var_after_op(anchor_op=op, old_var=h, new_var=new_h_all)
        if cs.op != op:
            ops_to_remove.append(cs.op)
    else:
        raise ValueError("tf lstm op {} not supported. Only {} supported".format(op.op_type, SUPPORTED_TF_LSTM_OPS))

    block.remove_ops(ops_to_remove)
    return True

    Inputs:

    op (Operation): op.op_type must be 'tf_lstm_block_cell' or `tf_lstm_block`

    Returns:

    True if op can be represented by mb.lstm op in SSA. False otherwise
    
coremltools/converters/mil/frontend/tensorflow/ssa_passes/tf_lstm_to_core_lstm.py