code: def _flip_sequence(sequences: Array, seq_lengths: Array) -> Array:
      max_steps = sequences.shape[0]
      roll_amounts = max_steps - seq_lengths
      # roll initially puts padding at the front so when the sequence is reversed
      # (via [::-1]) the padding stays at the end
      return jax.vmap(
          partial(jnp.roll, axis=0), in_axes=(1, 0), out_axes=1
      )(sequences, roll_amounts)[::-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 14 for calling a global function with example usage: _flip_sequence(sequences, seq_lengths) and returns: jax