(f, init, xs, length=None, reverse=False)
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.
Scan a function over leading array axes while carrying along state.
Scan a function over leading array axes while carrying along state .
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))
def scan ( f , init , xs , length = None , reverse = False ) : 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 ) )