Broadcast probability functions with value and size#8105
Broadcast probability functions with value and size#8105ricardoV94 wants to merge 5 commits intopymc-devs:mainfrom
Conversation
ricardoV94
force-pushed
the
xrv_bcast
branch
from
cb9f326 to
cee1def
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #8105 +/- ##
==========================================
+ Coverage 90.89% 90.92% +0.02%
==========================================
Files 123 123
Lines 19501 19558 +57
==========================================
+ Hits 17726 17783 +57
Misses 1775 1775
🚀 New features to boost your workflow:
|
ricardoV94
force-pushed
the
xrv_bcast
branch
from
cee1def to
33f2265
Compare
| pytestmark = pytest.mark.filterwarnings( | ||
| "error", | ||
| "ignore::numba.core.errors.NumbaPerformanceWarning", | ||
| "ignore:create_index_for_new_dim:UserWarning", |
There was a problem hiding this comment.
already fixed upstream, but we haven't released yet
ricardoV94
changed the title
ricardoV94
force-pushed
the
xrv_bcast
branch
4 times, most recently
from
94d7d75 to
0cc66d5
Compare
ricardoV94
force-pushed
the
xrv_bcast
branch
from
d4f5a9b to
f8d5790
Compare
| logcdf = pt.switch( | ||
| pt.lt(backward_value, 0), | ||
| pt.log(pt.exp(logcdf_zero) - pt.exp(logcdf)), | ||
| logdiffexp(logcdf_zero, logcdf), |
There was a problem hiding this comment.
With the broadcasted logp, this started triggering pymc-devs/pytensor#1883
So I fixed it in our math helper to go around the bug, and avoid having to wait for a pytensor release
There was a problem hiding this comment.
add a TODO that links to the pytensor issue and says to remove this once its fixed upstream
There was a problem hiding this comment.
We have this helper in pymc.math anyway, so no need to revert anything? It just starts already with the optimized form
| - arviz>=0.13.0 | ||
| - blas | ||
| - cachetools>=4.2.1 | ||
| - cachetools>=4.2.1,<7 |
| if size_idx is not None: | ||
| size = dist_params[size_idx] | ||
| if params_idxs is not None: | ||
| dist_params = [dist_params[i] for i in params_idxs] |
There was a problem hiding this comment.
Refactor this out to a helper, it's repeated 4x
|
|
||
| def test_censored_categorical(self): | ||
| cat = pm.Categorical.dist([0.1, 0.2, 0.2, 0.3, 0.2], shape=(5,)) | ||
| cat = pm.Categorical.dist([0.1, 0.2, 0.2, 0.3, 0.2], shape=()) |
There was a problem hiding this comment.
Why change the shapes in these tests?
There was a problem hiding this comment.
Oh I think I get it, because previously you had to say 5 to match the incoming parameter shape?
There was a problem hiding this comment.
No the previous was just wrong but because it was ignored nothing happened.
|
|
||
| def _to_xtensor(var, op: MeasurableXTensorFromTensor): | ||
| def _to_tensor(op: MeasurableXTensorFromTensor, value: XTensorVariable) -> TensorVariable: | ||
| # Align dims thate are shared between value and op to the right |
There was a problem hiding this comment.
typo: that (can't suggest im in commit-by-commit mode)
| missing_axis = [ | ||
| i for i, dim in enumerate(op.dims, start=n_value_unique_dims) if dim not in value_dims_set | ||
| ] | ||
| return pt_expand_dims(value.values, axis=missing_axis) |
There was a problem hiding this comment.
It seems lke we ought to have a helper that does this on the pytensor side
| def _to_xtensor( | ||
| op: MeasurableXTensorFromTensor, value: XTensorVariable, var: TensorVariable | ||
| ) -> XTensorVariable: | ||
| value_unique_dims = [dim for dim in value.dims if dim not in op.dims] |
There was a problem hiding this comment.
Name could be better, because dims are always unique by construction. value_only_dims?
2 participants
Since our syntax is
foo(rv, value), we should respect the shape of the rv. We were only respecting the one implied by the parameters, but not size.pm.logp(pm.Normal.dist(size=(10, 3)), 0))would return a single item tensorThis PR also tests we are allowed to further broadcast with value.
pm.logp(pm.Normal.dist(size=(10, 3)), zeros((3, 10, 3))Impl-wise I think it's cleaner to not include size in the definitions themselves as they are our construct.
This should also be compatible out of the box when we replace ours by those in pytensor-distributions.