[ot][spam][crazy][data] transformer model 'attention' improvement

[k]

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.