Vectorize ARProcess()

[dylanhmorris]

This will allow it to play better with numpyro.plate (still not seamlessly, see here), among other things

• This follows the convention that whenever a multi-dim array is a timeseries, time is the first dimension. One reason for that convention (@sbidari maybe worth noting in Document the convention that the first axis is time whenever a multi-dimensional array is a timeseries #443): By default, numpyro.scan builds by augmenting the first dimension. So since we most often scan over time, this convention will make scans over time have sensible default output.
• Adds a bunch of tests for correctness
• Adds informative error messages when shapes cannot be broadcast properly
• Adds tests for the errors.
• Note for PRs such as Document the convention that the first axis is time whenever a multi-dimensional array is a timeseries #443: clearest way to take dot products along only certain axes of a multidimensional array of unknown dimensions appears to be via jnp.einsum. I think we should use it wherever that's the task Vectorize ARProcess() #439 (comment)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.03%. Comparing base (bfaa0c0) to head (4c2561c).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #439      +/-   ##
==========================================
+ Coverage   93.96%   94.03%   +0.06%     
==========================================
  Files          42       42              
  Lines         978      989      +11     
==========================================
+ Hits          919      930      +11     
  Misses         59       59              

Flag	Coverage Δ
unittests	94.03% <100.00%> (+0.06%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[damonbayer]

Since one of the intended use cases of this extension is compatibility with numpyro.plate I wonder if we should do some tests with plates (and more complicated setups like nested plates. I'm not immediately sure that all this will work).

[dylanhmorris]

Since one of the intended use cases of this extension is compatibility with numpyro.plate I wonder if we should do some tests with plates (and more complicated setups like nested plates. I'm not immediately sure that all this will work).

Good point will add

[damonbayer]

Good point will add

I might even suggest that plate be the only way we do multi-dimensional testing.

[dylanhmorris]

Good point will add

I might even suggest that plate be the only way we do multi-dimensional testing.

I wouldn't go that far; we want vectorizations / multi-dim behavior to be as expected outside of plate contexts too, no?

[AFg6K7h4fhy2]

Confirming that I have seen all changes made in this PR and (mostly) have absorbed the updates.