Repository: pymc-devs/pymc
Pull request: Implemented logprob for SpecifyShape and CheckandRaise
<!-- !! Thank your for opening a PR !! -->

**What is this PR about?**
This PR closes [6352](https://github.com/pymc-devs/pymc/issues/6352).
I have implemented ```logprob``` for ```SpecifyShape``` and ```CheckandRaise``` and also included a rewrite that converts a SpecifyShape or ```CheckandRaise``` of a ```MeasurableVariable``` into a ```MeasurableSpecifyShape/Assert Op```.
I have also modified the ```__init__``` file to be consistent with the modifications.

**Checklist**
+ [x] Explain important implementation details 👆
+ [x] Make sure that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html).
+ [x] Link relevant issues (preferably in [nice commit messages](https://tbaggery.com/2008/04/19/a-note-about-git-commit-messages.html))
+ [x] Are the changes covered by tests and docstrings?
+ [x] Fill out the short summary sections 👇

## Major / Breaking Changes
- no

## New features
- Created a file to implement ```logprob``` for ```SpecifyShape``` by transfering ```specify_shape``` from ```rv``` to value and included a rewrite that converts a ```SpecifyShape``` of a ```MeasurableVariable``` into a ```MeasurableSpecifyShape Op```.
- Created a file to implement ```logprob``` for ```CheckandRaise```. If the assertion is true, the value variable will retain its value and if its false, a ```ValueError``` will be raised. Also have included a rewrite that converts a ```CheckandRaise``` of a ```MeasurableVariable``` into a ```MeasurableAssert Op```.

## Bugfixes
- no

## Documentation
- Documentation has been updated consistent with modifications

## Maintenance
- no


Problem statement:
Implement logprob for SpecifyShape and CheckAndRaise
### Description

The logp should be pretty simple:

For SpecifyShape:
```python
def logprob_specify_shape(op, values, inner_rv, *shapes, **kwargs):
  (value,) = values
  # transfer specify_shape from rv to value
  value = at.specify_shape(value, shapes)
  return logprob(inner_rv, value)
```

Similarly, for CheckAndRaise (e.g., `assert_op`), it would involve transfering the CheckAndRaise operation from the RV to the value variable

Hints from issue discussion:
I could take care of that, my question would be: Do is there some reason why we want to call the function with _specify_shape postfix? I assume that SpecifyShape is the class that needs this function, so it would be redundant to include it
I have no strong opinion about the name, just picked that to make the example readable. 

Better to check what names we give to other logprob functions in the logprob submodule and try to be consistent with that.

Note that we also need a rewrite that converts a SpecifyShape or Assert of a MeasurableVariable into a MeasurableSpecifyShape/Assert Op.

This might be a good template https://github.com/pymc-devs/pymc/blob/main/pymc/logprob/cumsum.py
Just to make sure, following the example from cumsum, I should create a file with a class extending SpecifyShape op, and within the class create the logprob_specify_shape function right? Then do accordingly with Assert, all within the logprob folder
> Just to make sure, following the example from cumsum, I should create a file with a class extending SpecifyShape op, and within the class create the logprob_specify_shape function right? Then do accordingly with Assert, all within the logprob folder

Yes :)
I think I might have overshot with taking this task so early in my learning process. If anyone would like to take it and finish, feel free to do so. I need to start with some easier tasks than pytensor functionality adding, as I don't fully understand nor know the library well enough to add this.
> I think I might have overshot with taking this task so early in my learning process. If anyone would like to take it and finish, feel free to do so. I need to start with some easier tasks than pytensor functionality adding, as I don't fully understand nor know the library well enough to add this.

That's totally fine. Thanks for giving it a try
Hey I would like to work on this issue.
Before starting, I have a question. I referred https://github.com/pymc-devs/pymc/blob/main/pymc/logprob/cumsum.py and from my understanding, here ```MeasurableCumsum``` is a subclass of ```CumOp``` defined as an extra_ops in Pytensor. I would like to know what class will our ```MeasurableSpecifyShape``` inherit from and if not what will we pass as argument to the decorator?
Basically, I would like to know what class would ```@node_rewriter([Op])``` contain in place of Op and if not will passing ```None``` be acceptable?
You should subclass the original Class, so Assert or SpecifyShape
I have a question. While transfering the Assert from the RV to the value variable, is using ```assert_op``` a possible solution ?
```
def logprob_assert(op, values, inner_rv, *assertion, **kwargs):
    (value,) = values
    # transfer assertion from rv to value
    value = at.assert_op(assertion, value)
    return logprob(inner_rv, value)
```
If the assertion is true, the value variable will retain its value and if its false, a ValueError will be raised.
Yes that's the idea, but you would want to use the same kind of `CheckAndRaise` `Op` that was used in the original graph. `CheckAndRaise` can be subclassed to raise specific error types with specific messages. `assert_op` is one such subclass that raises an `AssertionError`. 

The original issue should have been called "Implement logprob for SpecifyShape and CheckAndRaise". Renamed it now.

This is not more complicated, as you have access to the `Op` in the logprob function, so you can simply do `op(value, assertion)`

Commit message: Updated and rectified the checks in logprob and the corresponding tests suites

Patch to review:
diff --git a/pymc/logprob/__init__.py b/pymc/logprob/__init__.py
index 0b4bfc3385..f6ae51408c 100644
--- a/pymc/logprob/__init__.py
+++ b/pymc/logprob/__init__.py
@@ -40,6 +40,7 @@
 # Add rewrites to the DBs
 import pymc.logprob.censoring
 import pymc.logprob.cumsum
+import pymc.logprob.checks
 import pymc.logprob.mixture
 import pymc.logprob.scan
 import pymc.logprob.tensor
diff --git a/pymc/logprob/checks.py b/pymc/logprob/checks.py
new file mode 100644
index 0000000000..550d0e2d4e
--- /dev/null
+++ b/pymc/logprob/checks.py
@@ -0,0 +1,158 @@
+#   Copyright 2023 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+#   MIT License
+#
+#   Copyright (c) 2021-2022 aesara-devs
+#
+#   Permission is hereby granted, free of charge, to any person obtaining a copy
+#   of this software and associated documentation files (the "Software"), to deal
+#   in the Software without restriction, including without limitation the rights
+#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#   copies of the Software, and to permit persons to whom the Software is
+#   furnished to do so, subject to the following conditions:
+#
+#   The above copyright notice and this permission notice shall be included in all
+#   copies or substantial portions of the Software.
+#
+#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+#   SOFTWARE.
+
+from typing import List, Optional
+
+import pytensor.tensor as pt
+
+from pytensor.graph.rewriting.basic import node_rewriter
+from pytensor.raise_op import CheckAndRaise, ExceptionType
+from pytensor.tensor.shape import SpecifyShape
+
+from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
+from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+from pymc.logprob.utils import ignore_logprob
+
+
+class MeasurableSpecifyShape(SpecifyShape):
+    """A placeholder used to specify a log-likelihood for a specify-shape sub-graph."""
+
+
+MeasurableVariable.register(MeasurableSpecifyShape)
+
+
+@_logprob.register(MeasurableSpecifyShape)
+def logprob_specify_shape(op, values, inner_rv, *shapes, **kwargs):
+    (value,) = values
+    # transfer specify_shape from rv to value
+    value = pt.specify_shape(value, shapes)
+    return logprob(inner_rv, value)
+
+
+@node_rewriter([SpecifyShape])
+def find_measurable_specify_shapes(fgraph, node) -> Optional[List[MeasurableSpecifyShape]]:
+    r"""Finds `SpecifyShapeOp`\s for which a `logprob` can be computed."""
+
+    if isinstance(node.op, MeasurableSpecifyShape):
+        return None  # pragma: no cover
+
+    rv_map_feature: Optional[PreserveRVMappings] = getattr(fgraph, "preserve_rv_mappings", None)
+
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    rv = node.outputs[0]
+
+    base_rv, *shape = node.inputs
+
+    if not (
+        base_rv.owner
+        and isinstance(base_rv.owner.op, MeasurableVariable)
+        and base_rv not in rv_map_feature.rv_values
+    ):
+        return None  # pragma: no cover
+
+    new_op = MeasurableSpecifyShape()
+    # Make base_var unmeasurable
+    unmeasurable_base_rv = ignore_logprob(base_rv)
+    new_rv = new_op.make_node(unmeasurable_base_rv, *shape).default_output()
+    new_rv.name = rv.name
+
+    return [new_rv]
+
+
+measurable_ir_rewrites_db.register(
+    "find_measurable_specify_shapes",
+    find_measurable_specify_shapes,
+    "basic",
+    "specify_shape",
+)
+
+
+class MeasurableCheckAndRaise(CheckAndRaise):
+    """A placeholder used to specify a log-likelihood for an assert sub-graph."""
+
+
+MeasurableVariable.register(MeasurableCheckAndRaise)
+
+
+@_logprob.register(MeasurableCheckAndRaise)
+def logprob_assert(op, values, inner_rv, *assertion, **kwargs):
+    (value,) = values
+    # transfer assertion from rv to value
+    value = op(assertion, value)
+    return logprob(inner_rv, value)
+
+
+@node_rewriter([CheckAndRaise])
+def find_measurable_asserts(fgraph, node) -> Optional[List[MeasurableCheckAndRaise]]:
+    r"""Finds `AssertOp`\s for which a `logprob` can be computed."""
+
+    if isinstance(node.op, MeasurableCheckAndRaise):
+        return None  # pragma: no cover
+
+    rv_map_feature: Optional[PreserveRVMappings] = getattr(fgraph, "preserve_rv_mappings", None)
+
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    rv = node.outputs[0]
+
+    base_rv, *conds = node.inputs
+
+    if not (
+        base_rv.owner
+        and isinstance(base_rv.owner.op, MeasurableVariable)
+        and base_rv not in rv_map_feature.rv_values
+    ):
+        return None  # pragma: no cover
+
+    exception_type = ExceptionType()
+    new_op = MeasurableCheckAndRaise(exc_type=exception_type)
+    # Make base_var unmeasurable
+    unmeasurable_base_rv = ignore_logprob(base_rv)
+    new_rv = new_op.make_node(unmeasurable_base_rv, *conds).default_output()
+    new_rv.name = rv.name
+
+    return [new_rv]
+
+
+measurable_ir_rewrites_db.register(
+    "find_measurable_asserts",
+    find_measurable_asserts,
+    "basic",
+    "assert",
+)
diff --git a/tests/logprob/test_checks.py b/tests/logprob/test_checks.py
new file mode 100644
index 0000000000..6c890ed993
--- /dev/null
+++ b/tests/logprob/test_checks.py
@@ -0,0 +1,107 @@
+#   Copyright 2023 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+#   MIT License
+#
+#   Copyright (c) 2021-2022 aesara-devs
+#
+#   Permission is hereby granted, free of charge, to any person obtaining a copy
+#   of this software and associated documentation files (the "Software"), to deal
+#   in the Software without restriction, including without limitation the rights
+#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#   copies of the Software, and to permit persons to whom the Software is
+#   furnished to do so, subject to the following conditions:
+#
+#   The above copyright notice and this permission notice shall be included in all
+#   copies or substantial portions of the Software.
+#
+#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+#   SOFTWARE.
+
+import inspect
+import os
+import re
+import sys
+
+import numpy as np
+import pytensor
+import pytensor.tensor as pt
+import pytest
+
+from pytensor.raise_op import Assert
+from scipy import stats
+
+from pymc.distributions import Dirichlet
+from pymc.logprob.joint_logprob import factorized_joint_logprob
+
+currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
+parentdir = os.path.dirname(currentdir)
+sys.path.insert(0, parentdir)
+from distributions.test_multivariate import dirichlet_logpdf
+
+
+def test_specify_shape_logprob():
+    # 1. Create graph using SpecifyShape
+    # Use symbolic last dimension, so that SpecifyShape is not useless
+    last_dim = pt.scalar(name="last_dim", dtype="int64")
+    x_base = Dirichlet.dist(pt.ones((last_dim,)), shape=(5, last_dim))
+    x_base.name = "x"
+    x_rv = pt.specify_shape(x_base, shape=(5, 3))
+    x_rv.name = "x"
+
+    # 2. Request logp
+    x_vv = x_rv.clone()
+    [x_logp] = factorized_joint_logprob({x_rv: x_vv}).values()
+
+    # 3. Test logp
+    x_logp_fn = pytensor.function([last_dim, x_vv], x_logp)
+
+    # 3.1 Test valid logp
+    x_vv_test = stats.dirichlet(np.ones((3,))).rvs(size=(5,))
+    np.testing.assert_array_almost_equal(
+        x_logp_fn(last_dim=3, x=x_vv_test),
+        dirichlet_logpdf(x_vv_test, np.ones((3,))),
+    )
+
+    # 3.2 Test shape error
+    x_vv_test_invalid = stats.dirichlet(np.ones((1,))).rvs(size=(5,))
+    with pytest.raises(TypeError, match=re.escape("not compatible with the data's ((5, 1))")):
+        x_logp_fn(last_dim=1, x=x_vv_test_invalid)
+
+
+def test_assert_logprob():
+    rv = pt.random.normal()
+    assert_op = Assert("Test assert")
+    # Example: Add assert that rv must be positive
+    assert_rv = assert_op(rv > 0, rv)
+    assert_rv.name = "assert_rv"
+
+    assert_vv = assert_rv.clone()
+    assert_logp = factorized_joint_logprob({assert_rv: assert_vv})[assert_vv]
+
+    # Check valid value is correct and doesn't raise
+    # Since here the value to the rv satisfies the condition, no error is raised.
+    valid_value = 3.0
+    with pytest.raises(AssertionError, match="Test assert"):
+        assert_logp.eval({assert_vv: valid_value})
+
+    # Check invalid value
+    # Since here the value to the rv is negative, an exception is raised as the condition is not met
+    with pytest.raises(AssertionError, match="Test assert"):
+        assert_logp.eval({assert_vv: -5.0})


Code review comments:
diff_hunk: @@ -0,0 +1,107 @@
+#   Copyright 2023 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+#   MIT License
+#
+#   Copyright (c) 2021-2022 aesara-devs
+#
+#   Permission is hereby granted, free of charge, to any person obtaining a copy
+#   of this software and associated documentation files (the "Software"), to deal
+#   in the Software without restriction, including without limitation the rights
+#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#   copies of the Software, and to permit persons to whom the Software is
+#   furnished to do so, subject to the following conditions:
+#
+#   The above copyright notice and this permission notice shall be included in all
+#   copies or substantial portions of the Software.
+#
+#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+#   SOFTWARE.
+
+import inspect
+import os
+import re
+import sys
+
+import numpy as np
+import pytensor
+import pytensor.tensor as pt
+import pytest
+
+from pytensor.raise_op import Assert
+from scipy import stats
+
+from pymc.distributions import Dirichlet
+from pymc.logprob.joint_logprob import factorized_joint_logprob
+
+currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
+parentdir = os.path.dirname(currentdir)
+sys.path.insert(0, parentdir)
path: tests/logprob/test_checks.py
text: @user1:
This should be removed. For running the test locally, just make sure you are running from the root folder I think
```suggestion
```