code: def _select_last_carry(carry_seq: Array, seq_lengths: Array):
      return carry_seq[seq_lengths - 1, jnp.arange(carry_seq.shape[1])]
  
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 13 for calling a global function with example usage: _select_last_carry(carry_seq, seq_lengths) and returns: carry_seq