def scan(f, init, xs, length=None, reverse=False):
  """Scan a function over leading array axes while carrying along state.

  See the docstring of `jax.lax.scan`
  (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) for
  details.

  Args:
    f: a Python function to be scanned of type ``c -> a -> (c, b)``, meaning
      that ``f`` accepts two arguments where the first is a value of the loop
      carry and the second is a slice of ``xs`` along its leading axis, and that
      ``f`` returns a pair where the first element represents a new value for
      the loop carry and the second represents a slice of the output. Note that
      the input and output carry must have the same dtype.
    init: an initial loop carry value of type ``c``, which can be a scalar,
      array, or any pytree (nested Python tuple/list/dict) thereof, representing
      the initial loop carry value. This value must have the same structure as
      the first element of the pair returned by ``f``.
    xs: the value of type ``[a]`` over which to scan along the leading axis,
      where ``[a]`` can be an array or any pytree (nested Python
      tuple/list/dict) thereof with consistent leading axis sizes.
    length: optional integer specifying the number of loop iterations, which
      must agree with the sizes of leading axes of the arrays in ``xs`` (but can
      be used to perform scans where no input ``xs`` are needed).
    reverse: optional boolean specifying whether to run the scan iteration
      forward (the default) or in reverse, equivalent to reversing the leading
      axes of the arrays in both ``xs`` and in ``ys``.

  Returns:
    A pair of type ``(c, [b])`` where the first element represents the final
    loop carry value and the second element represents the stacked outputs of
    the second output of ``f`` when scanned over the leading axis of the inputs.
  """
  init, xs = tf.nest.map_structure(
      lambda x: tf_np.asarray(x) if x is not None else None, (init, xs))
  if length is not None:
    length = int(length)
  def get_length(x):
    if x is None:
      return None
    if x.shape.rank == 0:
      raise ValueError("Some array in `xs` doesn't have a leading dimension")
    return x.shape[0]
  lengths = tf.nest.flatten(tf.nest.map_structure(get_length, xs))
  for l in lengths:
    if l is not None:
      if length is None:
        length = l
      elif length != l:
        raise ValueError("There are two different leading-dimension lengths: "
                         f"{length} and {l}")
  if length is None:
    raise ValueError(
        "Can't determine length. Please set the `length` argument.")
  xs_ta = tf.nest.map_structure(
      lambda t: (tf.TensorArray(t.dtype, size=length, dynamic_size=False)  # pylint: disable=g-long-lambda
                 .unstack(t) if t is not None else None),
      xs)
  # tf.while_loop doesn't allow None in loop_vars, so we mask them.
  is_init_none = tf.nest.map_structure(lambda x: x is None, init)
  def to_safe(carry):
    return tf.nest.map_structure(
        lambda x, is_none: tf.zeros([]) if is_none else x, carry, is_init_none)
  def from_safe(safe_carry):
    return tf.nest.map_structure(
        lambda x, is_none: None if is_none else x, safe_carry, is_init_none)
  def body(i, safe_carry, ys_ta):
    carry = from_safe(safe_carry)
    if reverse:
      i_ = length - 1 - i
    else:
      i_ = i
    xs = tf.nest.map_structure(
        lambda x_ta: x_ta.read(i_) if x_ta is not None else None, xs_ta)
    carry, ys = f(*_tf_to_np((carry, xs)))
    ys_ta = tf.nest.map_structure(
        lambda y_ta, y: (y_ta.write(i_, y) if y is not None else y_ta),
        ys_ta, ys)
    i = i + 1
    safe_carry = to_safe(carry)
    return i, safe_carry, ys_ta
  xs_spec = tf.nest.map_structure(
      lambda t: tf.TensorSpec(t.shape[1:], t.dtype) if t is not None else None,
      xs)
  _, ys_spec = eval_on_shapes(f)(init, xs_spec)
  # ys_ta can't contain None because tf.while_loop doesn't allow None in
  # loop_vars.
  ys_ta = tf.nest.map_structure(
      lambda y: tf.TensorArray(y.dtype if y is not None else tf.float32,  # pylint: disable=g-long-lambda
                               size=length, dynamic_size=False),
      ys_spec)
  safe_init = to_safe(init)
  _, safe_carry, ys_ta = tf.while_loop(
      lambda i, *_: i < length, body, (0, safe_init, ys_ta),
      maximum_iterations=length)
  carry = from_safe(safe_carry)
  def _stack(a, spec):
    if spec is None:
      return None
    a = a.stack()
    a.set_shape((length,) + a.shape[1:])
    return a
  ys = tf.nest.map_structure(_stack, ys_ta, ys_spec)
  return _tf_to_np((carry, ys))
Scan a function over leading array axes while carrying along state.

  See the docstring of `jax.lax.scan`
  (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) for
  details.

  Args:
    f: a Python function to be scanned of type ``c -> a -> (c, b)``, meaning
      that ``f`` accepts two arguments where the first is a value of the loop
      carry and the second is a slice of ``xs`` along its leading axis, and that
      ``f`` returns a pair where the first element represents a new value for
      the loop carry and the second represents a slice of the output. Note that
      the input and output carry must have the same dtype.
    init: an initial loop carry value of type ``c``, which can be a scalar,
      array, or any pytree (nested Python tuple/list/dict) thereof, representing
      the initial loop carry value. This value must have the same structure as
      the first element of the pair returned by ``f``.
    xs: the value of type ``[a]`` over which to scan along the leading axis,
      where ``[a]`` can be an array or any pytree (nested Python
      tuple/list/dict) thereof with consistent leading axis sizes.
    length: optional integer specifying the number of loop iterations, which
      must agree with the sizes of leading axes of the arrays in ``xs`` (but can
      be used to perform scans where no input ``xs`` are needed).
    reverse: optional boolean specifying whether to run the scan iteration
      forward (the default) or in reverse, equivalent to reversing the leading
      axes of the arrays in both ``xs`` and in ``ys``.

  Returns:
    A pair of type ``(c, [b])`` where the first element represents the final
    loop carry value and the second element represents the stacked outputs of
    the second output of ``f`` when scanned over the leading axis of the inputs.
  
trax/tf_numpy/extensions/extensions.py