add `step` function to `S5SSM` for autoregressive generation

[ezhang94]

Primary changes

tests:

• verify equivalence of apply_ssm using parallel scan to simple for loop implementation
• verify equivalence of ar iteration through S5SSM.step to single S5SSM.__call__

additionally:

• merged apply_ssm_with_feedthrough into apply_ssm and return both the last state and the output sequence
• updated S5SSM docstring
• for C_init='complex_normal', explicitly cast C as a complex dtype

Other changes

Some minor modifications to simple_lm.py and ICLDataset

• raise Exceptions if layer name is not correctly requested
• add aux_rng_stream=['dropout'] at model level (not fully standardized yet)
  • to train with dropout, users should just have to specify dropout rate, and not have to mess with setting up RNG streams
• expose dataloader arguments for ICLDataset so that we can pass in custom collate_fn
  • e.g. collate_fn handles casting of arrays to jax.Array and shaping (batch_sz,) -> (n_devices, batch_sz_per_device,)

[ezhang94]

Let's "Rebase and merge" instead of "Squash and merge" because I have some minor commits here that creeped in and are unrelated to the AR generation.

[ezhang94]

Adding @ekellbuch for visibility of changes (and of course any review)