lax.associative_scan: fix docstring examples (#4172)
lax.associative_scan: fix docstring examples (#4172)

* lax.associative_scan: fix docstring examples
* add verbiage from #3583
google/jax,google/jax,google/jax,tensorflow/probability,google/jax,tensorflow/probability