21:  def chunk_scanner(idx):

idx comes from elements of jnp.arange() on line 31, which generates a
vector of ascending integers starting at 0.

Reply via email to