Reworks InitializeInfectionsExponentialGrowth to work with numpyro.plate

[damonbayer]

No description provided.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 93.51%. Comparing base (adf5aaf) to head (dc85fae).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #432      +/-   ##
==========================================
- Coverage   93.53%   93.51%   -0.02%     
==========================================
  Files          41       41              
  Lines        1021     1018       -3     
==========================================
- Hits          955      952       -3     
  Misses         66       66              

Flag	Coverage Δ
unittests	93.51% <100.00%> (-0.02%)	⬇️

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]

@dylanhmorris @sbidari docs are not rendering. The problem is that the shape of the output of the initialization method for a non-plated context has changed from (x,) to (x,1) (in a plated context is will be (x, n_plates)). Probably worth discussing a general plan for handling this type of change throughout PyRenew before I attempt a fix.

[dylanhmorris]

@dylanhmorris @sbidari docs are not rendering. The problem is that the shape of the output of the initialization method for a non-plated context has changed from (x,) to (x,1) (in a plated context is will be (x, n_plates)). Probably worth discussing a general plan for handling this type of change throughout PyRenew before I attempt a fix.

Is there a reason not to use jnp.squeeze in contexts like this?

[damonbayer]

@dylanhmorris that could be a good fix for the current demo models, but I think we need a more thorough exploration of multiple plated components working together.

[dylanhmorris]

Given the content of the change, could we have tests both for the case in which rate is a scalar and for the case in which it is an array?

[dylanhmorris]

LGTM. Tests could probably be DRY-ified via pytest.mark.parametrize but that can be a separate PR.

[damonbayer]

LGTM. Tests could probably be DRY-ified via pytest.mark.parametrize but that can be a separate PR.

I thought about that, but the logic to determine the manual computations seemed nearly as verbose as the fully manual solution, while being less readable.