DOC: increase clarity of `zerosum_axes` description for `py.ZeroSumNormal`

[drbenvincent]

Issue with current documentation:

Tagging @aseyboldt and @AlexAndorra

The current docs for the zerosum_axes kwarg for ZeroSumNormal is a little unclear, or perhaps slightly ambiguous, and could lead to incorrect parsing by readers.

Currently:

zerosum_axes: int, defaults to 1
Number of axes along which the zero-sum constraint is enforced, starting from the rightmost position. Defaults to 1, i.e the rightmost axis.

When I read this initially I parsed that as "Number of the axis along which the zero-sum constraint is enforced". Based on that (incorrect) reading of the sentence I then passed 0 to the zerosum_axes, which then resulted in an error.

I then thought that this might be an off by 1 error in how the zerosum_axes was processed. But then I got input from others that finally corrected my understanding.

The way how it works, by choosing the number of axes, counting from the last forwards, means that it is not possible to use this kwarg to state that you only want the first (of multiple) axes to have the zero sum constraint.

There is a relevant warning on the page at the moment:

zerosum_axes has to be > 0. If you want the behavior of zerosum_axes = 0, just use pm.Normal.

But I'm not convinced that's as accurate or as helpful as it could be. In my particular use case, I have a 2D ZeroSumNormal and want the sum to zero constraint to be applied only to the first (0th) dimension.

Idea or request for content:

• The description of the kwarg could have clearer wording. It could be made more clear that it operates by saying how many of the axes (from the end, backwards) will have the zero sum constraint. There is a relevant warning on the page
• There could be additional explanation that if a user wants to apply the zero sum constraint to the first dimension, or the first 2 of 3 dimensions etc, then the user will (currently) have to change the order of the dims and then transpose, or use moveaxis to get the dims in the order that a user needs.
• A bonus would be to include an example of doing that.
• An even bigger bonus would be to change the API to allow a user more freedom to specify the axes that the zero sum constraint applies to, without having to manually move the dimensions around. One example would be to allow a user to list a set of dimension indexes. Another idea would be to add a kwarg that toggles whether zerosum_axes operates from the first axis forward, or the final axis backwards.

[ricardoV94]

We could rename it to n_zerosum_axis to emphasize it's a quantity not a coordinate. About the lack of flexibility, the reason has to do with the RandomVariable API requiring (so far) support dimensions to be on the rightmost position:

https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/dimensionality.html#id1

Other PyMC components (e.g. shape argument, mixtures, rewrites) assume this is the case.

[drbenvincent]

Better yet, the plural n_zerosum_axes

[ricardoV94]

Better yet, the plural n_zerosum_axes

Yeah that's what I meant, I always mix the two. Data datum :P

[michaelraczycki]

Is this still waiting to be updated? I could take that

[drbenvincent]

As far as I know, nobody has tackled this issue. Feel free...

[michaelraczycki]

To sum up, because above you've thrown quite a lot of ideas:
The goal here is to change the name of the parameter, and for the old one include a deprecation warning, and ideally add an example of some usage? I could try to work on the last point about alerting the existing API, but I'm not sure if I understand it correctly, so far the zerosum_axes allows only to specify one single axes to go along, but you'd like to make it possible to specify multiple axes?

[ricardoV94]

zerosum_axes allows you to specify how many axis should be jointly zero-summed. This is not obvious from the name, hence the suggestion to rename it to n_zerosum_axes. No behavior has to be changed, just the deprecation warning like you said.

Right now setting zerosum_axes=2 means that the last two axes of the distribution will be jointly zero-summed, the same will be the case with n_zerosum_axes=2. By default only the last axis is zero-summed, which is what the docstrings try to say.

[michaelraczycki]

What do you think about the 'check_zerosum_axes' ? It directly refers to the param we're changing, but on the other hand to also provide backward support for old function names would introduce a lot of mess into the code. So we could just leave it as it is for now

[ricardoV94]

That's an internal function, you can rename it just fine

[ricardoV94]

Closed via #6522