Fix: _unstack_scan_params breaks after flax deserialization (from_bytes returns numpy arrays) (#1)

baseline: perf/mfu-optim branch from GitHub

initial commit