def prepare_lstm_weights(lstm, kernel, recurrent_kernel, bias, device):
    """Copies kernel and recurrent kernel weights in the Pytorch format
    We split the kernel and recurrent kernel weights, create associated
    torch tensors adapted to be in line with the Cudnn optimization.
    After we have copied the weights, we ensure the paramters are on
    the same device and memory layout is optimized for Cudnn.

    """

    lstm = lstm.to(device)
    hidden_size = lstm.hidden_size

    # Convert gates from Keras [i,f,c,o] to PyTorch [i,f,g,o]
    i_k, f_k, c_k, o_k = np.split(kernel, 4, axis=1)
    weight_ih_data = np.concatenate([i_k, f_k, c_k, o_k], axis=1).T

    i_r, f_r, c_r, o_r = np.split(recurrent_kernel, 4, axis=1)
    weight_hh_data = np.concatenate([i_r, f_r, c_r, o_r], axis=1).T

    if bias is not None:
        # Split Keras combined bias into input and hidden biases
        bias_ih_data = convert_to_tensor(bias, dtype="float32")
        bias_hh_data = torch.zeros_like(bias_ih_data)

    else:
        bias_ih_data = torch.zeros(4 * hidden_size, device=device)
        bias_hh_data = torch.zeros(4 * hidden_size, device=device)

    # Create PyTorch tensors for weights
    weight_ih = convert_to_tensor(weight_ih_data, dtype="float32").contiguous()
    weight_hh = convert_to_tensor(weight_hh_data, dtype="float32").contiguous()
    bias_ih = convert_to_tensor(bias_ih_data, dtype="float32").contiguous()
    bias_hh = convert_to_tensor(bias_hh_data, dtype="float32").contiguous()

    # Ensure the weights are all on the same device
    weight_ih = weight_ih.to(device)
    weight_hh = weight_hh.to(device)
    bias_ih = bias_ih.to(device)
    bias_hh = bias_hh.to(device)

    # Copy Keras weights into Torch's flat weights
    with torch.no_grad():
        lstm.weight_ih_l0.copy_(weight_ih)
        lstm.weight_hh_l0.copy_(weight_hh)
        lstm.bias_ih_l0.copy_(bias_ih)
        lstm.bias_hh_l0.copy_(bias_hh)

    # Optimize the layout
    lstm.flatten_parameters()

    # After prepare_lstm_weights:
    # Force all LSTM parameters to be on the correct device
    for param in lstm.parameters():
        if param.device != device:
            param.data = param.data.to(device)
Copies kernel and recurrent kernel weights in the Pytorch format
    We split the kernel and recurrent kernel weights, create associated
    torch tensors adapted to be in line with the Cudnn optimization.
    After we have copied the weights, we ensure the paramters are on
    the same device and memory layout is optimized for Cudnn.

    
keras/src/backend/torch/rnn.py