code: def _select_last_carry(sequence: A, seq_lengths: jnp.ndarray) -> A:
      last_idx = seq_lengths - 1
  
      def _slice_array(x: jnp.ndarray):
          return x[last_idx, jnp.arange(x.shape[1])]
  
      return jax.tree_map(_slice_array, sequence)
  
imports: from functools import partial
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import jax
  import numpy
  from absl import logging
  from jax import numpy
  from jax import random
  from typing_extensions import Protocol
  from flax.core import lift
  from flax.core.frozen_dict import FrozenDict
  from flax.linen import initializers, transforms
  from flax.linen.activation import sigmoid, tanh
  from flax.linen.dtypes import promote_dtype
  from flax.linen.linear import Conv, Dense, PrecisionLike, default_kernel_init
  from flax.linen.module import Module, compact, nowrap
label_desc: Use image node_id 1 for calling a global function with example usage: _select_last_carry(sequence, seq_lengths) and returns: jax, x