Fix dim lengths created from coords or ConstantData #5751

[LukeLB]

Closes #5751. Fixes dimensions created by add_coord by making dimension length a TensorConstant. I have also added a test for this using some code from #5181.

I'm setting as WIP as @canyon289 suggested this PR is required for other work to become unblocked and I'm having problems with some failing tests (which could just be due to me running things locally with an incorrect set-up?). So it's worth seeing if they fail on here and then look at fixes, which will be faster than me trying to figure things out on my own!

Depending on what your PR does, here are a few things you might want to address in the description:

• what are the (breaking) changes that this PR makes?
• important background, or details about the implementation
• are the changes—especially new features—covered by tests and docstrings?
• consider adding/updating relevant example notebooks
• right before it's ready to merge, mention the PR in the RELEASE-NOTES.md

[canyon289]

it's worth seeing if they fail on here and then look at fixes, which will be faster than me trying to figure things out on my own

Great strategy. Thanks for taking such a principled approach to development

It seems that there type being changed somewhere. Might be as simple as updating this test.
https://github.com/pymc-devs/pymc/runs/6435716258?check_suite_focus=true#step:7:80

[codecov]

Codecov Report

Merging #5763 (e97ad4e) into main (00d4372) will increase coverage by 0.00%.
The diff coverage is 93.10%.

❗ Current head e97ad4e differs from pull request most recent head abad0a4. Consider uploading reports for the commit abad0a4 to get more accurate results

@@           Coverage Diff           @@
##             main    #5763   +/-   ##
=======================================
  Coverage   89.27%   89.28%           
=======================================
  Files          74       74           
  Lines       13811    13822   +11     
=======================================
+ Hits        12330    12341   +11     
  Misses       1481     1481           

Impacted Files	Coverage Δ
pymc/data.py	81.63% <87.50%> (-1.99%)	⬇️
pymc/model.py	86.78% <100.00%> (+0.40%)	⬆️
pymc/sampling.py	88.42% <100.00%> (+0.05%)	⬆️
pymc/step_methods/hmc/base_hmc.py	89.68% <0.00%> (-0.80%)	⬇️
pymc/step_methods/hmc/nuts.py	97.50% <0.00%> (+2.50%)	⬆️

[LukeLB]

Failing test looks to be down to sampling error?

test_sample_posterior_predictive_after_set_data_with_coord
E       AssertionError: 
E       Not equal to tolerance rtol=1e-07, atol=0.1
E       
E       Mismatched elements: 2 / 2 (100%)
E       Max absolute difference: 0.13669218
E       Max relative difference: 0.02331315
E        x: array([5, 6])
E        y: array([4.89599 , 5.863308]) 

And it's highlighted in an outstanding issue #5739.

[LukeLB]

This test is failing because y is not the correct shape, y.eval().shape == (1, 1). Both feature and group are TensorConstants so the addition to add_coord is doing as intented with respect to type changes. I'm not sure I understand why this is still happening.....

def test_add_coord_fixed_dims():
    """
    test that dims does not cause resizing of y. See #5181.
    """
    with pm.Model(coords={"feature": [1], "group": ["A", "B"]}):
        x = pm.Normal("x", 0, 1, dims="feature")
        y = pm.Normal("y", x[..., None], 1, dims=("feature", "group"))
    assert x.eval().shape == (1,)
    assert y.eval().shape == (1, 2)

[michaelosthege]

Maybe do more asserts?

with pm.Model(coords={"feature": [1], "group": ["A", "B"]}):
    assert pmodel.dim_lengths["feature"].eval() == 1
    assert pmodel.dim_lengths["group"].eval() == 2

[LukeLB]

My thinking for that test was to make sure that size and dims can be passed at the same time and we get the expected shape out which will not have been possible before this type change in add_coord(). So maybe its worth having two tests, the one you've suggested, and the other we can use the case you highlight #5181 (comment) to test that size and dims can be passed at the same time. What do you think?

[michaelosthege]

When size/shape and dims are passed at the same time, the size/shape takes priority and dims are just tracked for labeling the dimensions in the final xarray objects.

That's why I suggested to only pass dims, and actually it's a little concerning that the test doesn't pass.

Maybe call add_coord directly while passing fixed. For the variables depending on it there should be no difference between aesara.shared() or at.constant() dim lengths.

with pm.Model() as pmodel:
    pmodel.add_coord("feature", [1], fixed=True/False)
    pmodel.add_coord("group", ["A", "B"], fixed=True/False)
    assert pmodel.dim_lengths["feature"].eval() == 1
    assert pmodel.dim_lengths["group"].eval() == 2
    x = pm.Normal("x", 0, 1, dims="feature")
    y = pm.Normal("y", x[..., None], 1, dims=("feature", "group"))
assert x.eval().shape == (1,)
assert y.eval().shape == (1, 2)

[LukeLB]

When I try that all cases of fixed = True/False lead to y.eval().shape == (1, 1).

[michaelosthege]

Oh, now I see why. The x[..., None] creates the RV as (1, 1) and because the dimensionality is as expected, there's no further resizing.

This works 👇 because then the RV is (1,) and repeated for the "group" dimension.

    # ...
    x = pm.Normal("x", 0, 1, dims="feature")
    y = pm.Normal("y", x, 1, dims=("group", "feature"))
assert x.eval().shape == (1,)
assert y.eval().shape == (2, 1)

[LukeLB]

I think I'm still misunderstanding something here. Wouldn't passing dims like that still be allowed prior to this PR? So I don't understand what it is testing with respect to this PR.

[michaelosthege]

Yes, that's already allowed, but where this test comes in is in combination with the changes from #5761.

Actually, it's a little unfortunate that these changes are on separate branches.
Your test case should (latest after adding a pm.set_data) close the coverage gaps from that other PR.

I'm inclined to merge #5761 even with those coverage gaps (I'll review it really carefully right now).
After that you could just git rebase origin/main, add two lines doing pm.set_data and assert the ShapeError or ShapeWarning and then we should be good to go 🤔

[michaelosthege]

@LukeLB the other PR was merged, so now you should be able to rebase and bring the coverage back up :)

[LukeLB]

@michaelosthege OK I think I understand. I'll rebase and commit your suggested changes.

[michaelosthege]

Can you add test cases for the ShapeWarning and ValueError?

With match="" one can make sure that the exception is the correct one. That's important because there are two different ShapeErrors that one could run into..

We should aim to cover all code paths from the diff here https://github.com/pymc-devs/pymc/pull/5761/files

[michaelosthege]

I can't make the change suggestion, but the whole block containing the ShapeWarning can be removed..

[michaelosthege]

Thanks @LukeLB !

This brings the coverage back up should lead to fewer shape problems long term.

[michaelosthege]

ah damn, I just realized that we forgot to actually use the new fixed kwarg.

I commented a change in one of the tests that should currently fail.
Let me know if you want me to take over - you put in a lot of effort already and I don't want to get demotivated by all this back and forth

[LukeLB]

ah damn, I just realized that we forgot to actually use the new fixed kwarg.

I commented a change in one of the tests that should currently fail. Let me know if you want me to take over - you put in a lot of effort already and I don't want to get demotivated by all this back and forth

Not at all! The longer it goes the more I learn :) Thanks for the offer though.

[LukeLB]

Why does the assert statement on line 758 fail? The ConstantData() function suggests that it returns a TensorConstant.

[michaelosthege]

Why does the assert statement on line 759 fail? Its created from coord so should be a TensorConstant but is Subtensor{int64}.0.

The Subtensor{int64}.0 is because from pm.ConstantData([1,2,3], dims="foo") it's effectively doing this: dim_lengths["foo"] = at.constant([1,2,3]).shape[0].

So what needs to be done is a change in the logic behind pm.ConstantData (or rather pm.Data) to determine the constant data's shape non-symbolically first, and then call add_coord(..., fixed=True).

And I should add that pm.MutableData([1,2,3], dims="bla") becomes: dim_lengths["bla"] = aesara.shared([1,2,3]).shape[0].

[LukeLB]

OK so I think I have a solution which I have pushed but I'm running into a wall with a failing test which I can't seem to fix. Unfortunatley I won't be able to work on this till after the weekend so if there is a time pressure on this then it will be worth someone else picking it up.

[michaelosthege]

Thanks @LukeLB, I will try to pick it up in the evening

[canyon289]

@LukeLB this weekend is fine! You've made great progress on this and don't want to take away joy of getting it merged!

[michaelosthege]

I rebased the branch and squashed the history into three logically separate commits:

1. ✔ Removing the ShapeWarning in favor of raising earlier
2. ✔ Adding optional coords to the pm.Data API, because we were taking coordinate values from the data variable, which does not work for >1 dimensions.
3. ✔ Addition of the fixed kwarg (the actual goal of this PR)

I will force-push them one by one so we can confirm that the first two commits don't create new problems.

[michaelosthege]

@LukeLB @canyon289 @ricardoV94 please review and note the new changes I made in the pm.Data API

[ricardoV94]

Looks good to me, left a small question and a suggestion about using the same keyword internally (fixed vs mutable)

[ricardoV94]

doesn't matter much but x.type.shape should have the same info

[michaelosthege]

I feel 👍 about merging as long as the tests go ✔.

[canyon289]

Me too!

[canyon289]

Congrats @LukeLB! Thank you for leading this PR

[LukeLB]

@canyon289 @michaelosthege @ricardoV94 thanks all!! It's been a pleasure :)