def METHOD_NAME(weights, x, h_0, c_0):
  W_ih, W_hh, b_ih, b_hh = rnn.unpack_lstm_weights(weights, input_size,
                                                  hidden_size, num_layers,
                                                  bidirectional)
  y_ref, h_n_ref, c_n_ref = rnn.lstm_ref(
      x,
      h_0,
      c_0,
      W_ih,
      W_hh,
      b_ih,
      b_hh,
      seq_lengths=seq_lengths,
      input_size=input_size,
      hidden_size=hidden_size,
      num_layers=num_layers,
      dropout=False,
      bidirectional=bidirectional)
  seq_length_mask = jnp.tile(jnp.arange(seq_len, dtype=jnp.int32)[None],
      [batch_size, 1]) < seq_lengths[:, None]
  loss = jnp.sum(jnp.where(seq_length_mask[..., None], y_ref, 0.))
  return loss, (y_ref, h_n_ref, c_n_ref)