Fix broadcasting via observed and dims

[ricardoV94]

Fixes #5993

This PR also removes the support of ellipsis in dims and shape. This feature was not very advertised so hopefully this won't cause too much pain.

These tests illustrate the bug fix

    def test_broadcast_by_dims(self):
        with pm.Model(coords={"broadcast_dim": range(3)}) as m:
            x = pm.Normal("x", mu=np.zeros((1,)), dims=("broadcast_dim",))
            assert x.eval().shape == (3,)

    def test_broadcast_by_observed(self):
        with pm.Model() as m:
            x = pm.Normal("x", mu=np.zeros((1,)), observed=np.zeros((3,)))
            assert x.eval().shape == (3,)

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• When shape is not provided, dims or observed are used to define the shape of a variable. If you want a variable that has dims or is observed to automatically resize when its inputs change, you must specify that explicitly via the shape argument. For example, pm.Normal("likelihood", mu=mu, sigma=sigma, observed=data, dims="data", shape=mu.shape)
• Remove support of Ellipsis (...) in shape and dims

Bugfixes / New features

• Fix bug where distribution shape would not be broadcasted by dims or observed
• Allow specifying dims on the fly from observed
• Do not show shape-related dependencies of RandomVariables in model_graph

Docs / Maintenance

• Update _make_nice_attr_error to suggest pm.draw instead of .eval()
• Add _make_nice_attr_error to SymbolicDistributions

[ricardoV94]

Started a discussion for removing dims on the fly: #6065

[ricardoV94]

I seem to have broken some things according to the failing test_density_scaling_with_generator. It seems that we support generators as observed data... which means we can only inquiry about data once (and then we must pass it around). Is this something we really need?

Seems to be used by Minibatch, cc @ferrine

[codecov]

Codecov Report

Merging #6063 (c9660b6) into main (fa5e441) will increase coverage by 0.01%.
The diff coverage is 96.66%.

@@            Coverage Diff             @@
##             main    #6063      +/-   ##
==========================================
+ Coverage   89.52%   89.53%   +0.01%     
==========================================
  Files          72       72              
  Lines       12944    12929      -15     
==========================================
- Hits        11588    11576      -12     
+ Misses       1356     1353       -3     

Impacted Files	Coverage Δ
pymc/distributions/distribution.py	91.02% <94.73%> (+0.18%)	⬆️
pymc/distributions/shape_utils.py	99.27% <100.00%> (+2.19%)	⬆️
pymc/model.py	88.23% <100.00%> (+0.03%)	⬆️
pymc/model_graph.py	78.94% <100.00%> (+0.63%)	⬆️
pymc/gp/gp.py	92.73% <0.00%> (-0.45%)	⬇️
pymc/distributions/multivariate.py	92.03% <0.00%> (+0.02%)	⬆️
pymc/step_methods/hmc/base_hmc.py	90.55% <0.00%> (+0.78%)	⬆️

[ricardoV94]

Tests are passing

[michaelosthege]

Tests are passing

👍 I will review this tonight!

[ricardoV94]

Tests are passing

+1 I will review this tonight!

I'll wait for it!

[michaelosthege]

Don't you want to broadcast in this test too?

Suggested change

	y = pm.Normal("y", observed=[0, 0, 0], dims="ddata")
	assert pmodel.RV_dims["y"] == ("ddata",)
	assert y.eval().shape == (3,)
	y = pm.Normal("y", observed=[[0, 0, 0]], dims=("dobs", "ddata"))
	assert pmodel.RV_dims["y"] == ("dobs", "ddata")
	assert y.eval().shape == (1, 3)

[ricardoV94]

Not sure, is this not already covered by test_define_dims_on_the_fly_from_observed? If it is the same I would like to keep it separate, in case we decide to drop support for dims on the fly.

[michaelosthege]

Signature violation: The previous dims: WeakDims parameter was already a sequence, but with dims: Dims the new signature allows for dims="bla".
This is not accounted for in the implementation; the cast(StrongDims, dims) is not safe yet.

[ricardoV94]

Removed the cast, now that the signature of shape_from_dims requires StrongDims

[ricardoV94]

@michaelosthege Do you see a difference between broadcasting by dims and by observed? Because all the same changes would have been needed if the tests had used dims instead of observed and then assumed you could just resize a variable indirectly by resizing it's arguments.

[michaelosthege]

@michaelosthege Do you see a difference between broadcasting by dims and by observed? Because all the same changes would have been needed if the tests had used dims instead of observed and then assumed you could just resize a variable indirectly by resizing it's arguments.

I see you point, but I think for most people the intuition is that shape, size and dims are an input of the RV, whereas the observed is something that the RV points to.
That's also how we draw it in pm.model_to_graphviz and it resembles the direction of the generative process.

That's probably why I don't like to broadcast automatically based on observed. (I found that strange and complicated when I worked on these implementations in the past..)

[ricardoV94]

The graphviz is a bit strange because we use 2 nodes (RV and Value) to represent what is done by a single one for unobserved RVs.

Anyway, I think:

1. The output of test_broadcast_by_observed would be very surprising before this PR. Also would be very surprising if we ignored extra dimensions implied by observed (which we don't)
2. We should make it clear (in the dimensionality notebook?) that in the generative graph, observed is used as a shortcut for shape=observed.shape unless dims or shape is provided. I hope this provides a good mental picture so users don't expect indirect resizing by default.
3. We should make it clear that the shape argument is more powerful than before. It can take arbitrary symbolic expressions. Hopefully it will become intuitive to specify shape=mu.shape or whatever as we start adding it to official examples.

I agree it's more cumbersome, but it does push towards being more explicit about shapes, which we are already doing anyway when we suggest users use dims.

I agree with making it a minor bigger (what's the name for intermediate?) release.

[michaelosthege]

I can't review super-thoroughly right now, but I trust that @ricardoV94 addressed my concerns.

I'm still not super happy about prioritizing observed.shape over the implied dimensionality, but we'll see how the users take it.
I know enough about PyMC so my own models don't blow up because of it ¯_(ツ)_/¯