code: def _extract_output(
      seq_lengths: Array, out
  ) -> tuple[tuple[Array, Array], Array]:
      _, ((hs, cs), seq_first_y) = out
      h_t = _select_last_carry(hs, seq_lengths)
      c_t = _select_last_carry(cs, seq_lengths)
  
      # [seq_len, batch]   [1, batch]             [seq_len, 1]
      mask = (
          seq_lengths[None]
          > jnp.arange(seq_first_y.shape[0], dtype=jnp.int32)[:, None]
      )
      # [batch, seq_len, hidden_size]
      seq_first_y = jnp.where(
          mask[..., None],  # [seq_len, batch, 1]
          seq_first_y,  # [seq_len, batch, hidden_size]
          0,
      )
      return (h_t, c_t), seq_first_y
  
imports: from functools import partial
  import math
  from typing import Any
  import jax
  import numpy
  from jax._src import core
  from jax.interpreters import mlir
  from jax.interpreters import xla
  from jax._src.custom_derivatives import custom_vjp
  from jax._src.typing import Array, Shape
  from jax._src.lax import lax
  import jax.numpy
label_desc: Use image node_id 12 for calling a global function with example usage: _extract_output(seq_lengths, out) and returns: h_t,  c_t,  seq_first_y