Suggestion: Replace _scan in batch.py with lax.scan

[mohamad-amin]

I was just reading through the file as I wanted to apply some modifications and I saw this function:

neural-tangents/neural_tangents/utils/batch.py

Line 126 in 3deb197

def _scan(f: Callable[[_Carry, _Input], Tuple[_Carry, _Output]],

And this comment:

"""Implements an unrolled version of scan.
Based on jax.lax.scan and has a similar API.
TODO(schsam): We introduce this function because lax.scan currently has a
higher peak memory usage than the unrolled version. We will aim to swap this
out for lax.scan when issue #1273 and related have been resolved.
"""

Which is fixed:

jax-ml/jax#1273