如何在Keras摘要中显示所有LSTM状态？

import keras
from keras import backend as K
from keras.models import Sequential
from keras.layers import LSTM

model = Sequential()
model.add(LSTM(units=128, return_sequences=True, input_shape=(timesteps, input_dim)))
model.compile(loss='mse', optimizer='adam')

def get_lstm_states(model, input_data):
    get_states = K.function([model.layers[0].input], [model.layers[0].output, model.layers[0].states[0], model.layers[0].states[1]])
    return get_states([input_data])[1:]

input_data = ...  # 输入数据
lstm_states = get_lstm_states(model, input_data)

print("LSTM状态1：", lstm_states[0])
print("LSTM状态2：", lstm_states[1])