（本貼文所列出的程式碼，皆以 colab 筆記本方式執行，[可由此下載]
(https://colab.research.google.com/drive/1tIu9KwFqp7dZ_vCLOZiQ0NK_Y0av0vGF?usp=sharing)
scan 算是最為複雜的 JAX 控制流程運算子，它其實和 TensorFlow 的 tf.scan() 很像，算是其簡化的版本。如果老頭的說明不夠詳盡，使得你不能完全領會 JAX scan 的用法的話，你不妨先去翻翻 tf.scan ，這會對你很有幫助。

scan() 運作的圖解

scan 的基本操作，是用掃描函式 f，依序以陣列元素為輸入值，產生掃描結果，最終將所有的掃描結果組成陣列回傳。

在掃描的過程中，f 也需要「傳承 carry」做為其輸入參數，並輸出下一個傳承，供掃描下一個陣列示素時使用。傳承的目的，是提供掃描程式執行時的「環境脈絡 context」，例如目前處理的陣列元素它的索引值等。

scan() 的語法及語義

f :

• 掃描函式 f ，它的輸入參數是「傳承 carry」和輸入陣列中的任一元素，它的輸出值為下一個傳承和掃描結果。
• 要注意的是，ｆ的輸出傳承值，必須維持和輸入傳承值一樣的維度及型態。

init :

• 掃描函式 f 的起始傳承值。

xs:

• 輸入陣列，陣列中的元素，是沿著陣列的第一個維度擺放的。（有些機器學習框架，元素沿著最後一個維度擺放)
• 元素本身，可以是一個陣列，或是其他 pytree 結構。（pytree 還沒有介紹，

length:

• 當 xs 不是 None 的時候，沒有必要使用 length。
• 當 xs 為 None 的時候，length 則用來指定 f 要掃描的次數。
• 老頭在稍後會舉一個例子。

reverse:

• 當為 True 時，由輸入陣列的尾端，反向掃描。
• 即定值是 False。

unroll:

• 迴圈展開的次數，既定值是 1，不做展開。
• 基本上展開的次數越多，所需的記憶體就越多，而程式執行就會越快。
• 設定 unroll 並不會影響最終的結果。

回傳值:

• 包含最終傳承值及輸出陣列的 tuple。

scan 的運作流程，可以用以下的 Python 程式段來說明：

def scan(f, init, xs, length=None):
  if xs is None:
    xs = [None] * length
  carry = init
  ys = []
  for x in xs:
    carry, y = f(carry, x)
    ys.append(y)
  return carry, np.stack(ys)

scan() 的使用範例

1. 計算陣列的累積和 (cumulative sum)
   

2. 計算 1 加到 10 的和

這個例子說明的輸入陣列 xs 為 none，藉由指定 length 的值來達到目的。

3. unroll 造成的效率差異

可以看出來，當設定較大的 unroll 值後，scan 的執行速度會略為加快。