Repository: pymc-devs/pymc

Issue:
BUG: KaTeX parse error on Jupyter Notebook when using underscore, etc.
### Describe the issue:

_I apologize if this was an intended feature implemented to reflect the demand of the PyMC community. But when I migrate my codebase from PyMC3 to PyMC5, this seems more like a bug to me; it didn't raise an error in PyMC 3.8._

## Issue
When trying to display (not print) a PyMC model on Jupyter Notebook, I received a "KaTeX parse error". Apparently, this is due to the underscore character, `_`, in the name.

### Reproduceable code example:

```python
import pymc as pm
with pm.Model() as model:
    new_a = pm.Normal('new_a')
display(model) # on jupyter notebook
```


### Error message:

```shell
ParseError: KaTeX parse error: Expected 'EOF', got '_' at position 53: … \text{new_̲a} &\sim & \ope…
```


### PyMC version information:

`pip install pymc==5.0.2`

### Context for the issue:

Right now, I could bypass this error by adding a slash, `\_`, e.g. `new\_a` instead of `new_a`. But this is less desirable. For instance, I often have many variables that I would like to define. To do that, I would loop over some dictionary, and declare the variable one by one, while automatically assign them names that match the keys in the dictionary. However, with this "bug", I would either have to avoid using underscores at all or pre-process the names by replacing `_` with `\_`.

### More
Because of this parse error, I thought maybe I could write LaTeX syntax in my name field, or use raw string? But I don't think it works. Below are the things that I have tried:
- `pm.Normal(r'new_a')` - I though raw string would "protect" the underscore from getting parse error, but it still raises the same error.
- `pm.Normal('\alpha')` - Nope.
- `pm.Normal('\\alpha')` - Nope.
- `pm.Normal('$\alpha$')` - Nope.
- `pm.Normal('$\\alpha$')` - Nope.


Discussion:
Dekermanjian: Hey @ricardoV94, I added a simplistic fix to address the underscores breaking the KaTeX output in VSCode and Jupyter. I tested in VSCode and Jupyter Lab. Are there any other environments I should test it out in? 

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7eee3bd4d..09bed88c4 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,6 +16,10 @@ repos:
     -   id: requirements-txt-fixer
         exclude: ^requirements-dev\.txt$
     -   id: trailing-whitespace
+- repo: https://github.com/citation-file-format/cffconvert
+  rev: 054bda51dbe278b3e86f27c890e3f3ac877d616c
+  hooks:
+    - id: validate-cff
 - repo: https://github.com/lucianopaz/head_of_apache
   rev: "0.0.3"
   hooks:
diff --git a/docs/source/learn/core_notebooks/GLM_linear.ipynb b/docs/source/learn/core_notebooks/GLM_linear.ipynb
index 89932575c..8991673c1 100644
--- a/docs/source/learn/core_notebooks/GLM_linear.ipynb
+++ b/docs/source/learn/core_notebooks/GLM_linear.ipynb
@@ -334,7 +334,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Much shorter, but this code does the exact same thing as the previous specification (you can change priors and everything else too if we wanted). `bambi` parses the `formulae` model string, adds random variables for each regressor (`Intercept` and slope `x` in this case), adds a likelihood (by default, a Normal is chosen), and all other variables (`sigma`). Finally, `bambi` then initializes the parameters to a good starting point by estimating a frequentist linear model using [statsmodels](http://statsmodels.sourceforge.net/devel/).\n",
+    "Much shorter, but this code does the exact same thing as the previous specification (you can change priors and everything else too if we wanted). `bambi` parses the `formulae` model string, adds random variables for each regressor (`Intercept` and slope `x` in this case), adds a likelihood (by default, a Normal is chosen), and all other variables (`sigma`). Finally, `bambi` then initializes the parameters to a good starting point by estimating a frequentist linear model using [statsmodels](http://statsmodels.sourceforge.net/).\n",
     "\n",
     "If you are not familiar with R's syntax, `'y ~ x'` specifies that we have an output variable `y` that we want to estimate as a linear function of `x`."
    ]
diff --git a/pymc/printing.py b/pymc/printing.py
index 56445ab9e..ef417f379 100644
--- a/pymc/printing.py
+++ b/pymc/printing.py
@@ -13,6 +13,8 @@
 #   limitations under the License.
 
 
+import re
+
 from functools import partial
 
 from pytensor.compile import SharedVariable
@@ -58,6 +60,7 @@ def str_for_dist(
     if "latex" in formatting:
         if print_name is not None:
             print_name = r"\text{" + _latex_escape(print_name.strip("$")) + "}"
+            print_name = _format_underscore(print_name)
 
         op_name = (
             dist.owner.op._print_name[1]
@@ -114,6 +117,7 @@ def str_for_model(model: Model, formatting: str = "plain", include_params: bool
     if not var_reprs:
         return ""
     if "latex" in formatting:
+        var_reprs = [_format_underscore(x) for x in var_reprs]
         var_reprs = [
             var_repr.replace(r"\sim", r"&\sim &").strip("$")
             for var_repr in var_reprs
@@ -295,3 +299,10 @@ try:
 except (ModuleNotFoundError, AttributeError):
     # no ipython shell
     pass
+
+
+def _format_underscore(variable: str) -> str:
+    """
+    Escapes all unescaped underscores in the variable name for LaTeX representation.
+    """
+    return re.sub(r"(?<!\\)_", r"\\_", variable)


Test patch:
diff --git a/tests/test_printing.py b/tests/test_printing.py
index 406032b12..832699c20 100644
--- a/tests/test_printing.py
+++ b/tests/test_printing.py
@@ -11,6 +11,7 @@
 #   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.
+
 import numpy as np
 
 from pytensor.tensor.random import normal
@@ -171,15 +172,15 @@ class TestMonolith(BaseTestStrAndLatexRepr):
                 r"$\text{mu} \sim \operatorname{Deterministic}(f(\text{beta},~\text{alpha}))$",
                 r"$\text{beta} \sim \operatorname{Normal}(0,~10)$",
                 r"$\text{Z} \sim \operatorname{MultivariateNormal}(f(),~f())$",
-                r"$\text{nb_with_p_n} \sim \operatorname{NegativeBinomial}(10,~\text{nbp})$",
+                r"$\text{nb\_with\_p\_n} \sim \operatorname{NegativeBinomial}(10,~\text{nbp})$",
                 r"$\text{zip} \sim \operatorname{MarginalMixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5))$",
                 r"$\text{w} \sim \operatorname{Dirichlet}(\text{<constant>})$",
                 (
-                    r"$\text{nested_mix} \sim \operatorname{MarginalMixture}(\text{w},"
+                    r"$\text{nested\_mix} \sim \operatorname{MarginalMixture}(\text{w},"
                     r"~\operatorname{MarginalMixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5)),"
                     r"~\operatorname{Censored}(\operatorname{Bernoulli}(0.5),~-1,~1))$"
                 ),
-                r"$\text{Y_obs} \sim \operatorname{Normal}(\text{mu},~\text{sigma})$",
+                r"$\text{Y\_obs} \sim \operatorname{Normal}(\text{mu},~\text{sigma})$",
                 r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",
                 r"$\text{pred} \sim \operatorname{Deterministic}(f(\text{<normal>}))",
             ],
@@ -189,11 +190,11 @@ class TestMonolith(BaseTestStrAndLatexRepr):
                 r"$\text{mu} \sim \operatorname{Deterministic}$",
                 r"$\text{beta} \sim \operatorname{Normal}$",
                 r"$\text{Z} \sim \operatorname{MultivariateNormal}$",
-                r"$\text{nb_with_p_n} \sim \operatorname{NegativeBinomial}$",
+                r"$\text{nb\_with\_p\_n} \sim \operatorname{NegativeBinomial}$",
                 r"$\text{zip} \sim \operatorname{MarginalMixture}$",
                 r"$\text{w} \sim \operatorname{Dirichlet}$",
-                r"$\text{nested_mix} \sim \operatorname{MarginalMixture}$",
-                r"$\text{Y_obs} \sim \operatorname{Normal}$",
+                r"$\text{nested\_mix} \sim \operatorname{MarginalMixture}$",
+                r"$\text{Y\_obs} \sim \operatorname{Normal}$",
                 r"$\text{pot} \sim \operatorname{Potential}$",
                 r"$\text{pred} \sim \operatorname{Deterministic}",
             ],
@@ -256,7 +257,7 @@ def test_model_latex_repr_three_levels_model():
         "$$",
         "\\begin{array}{rcl}",
         "\\text{mu} &\\sim & \\operatorname{Normal}(0,~5)\\\\\\text{sigma} &\\sim & "
-        "\\operatorname{HalfCauchy}(0,~2.5)\\\\\\text{censored_normal} &\\sim & "
+        "\\operatorname{HalfCauchy}(0,~2.5)\\\\\\text{censored\\_normal} &\\sim & "
         "\\operatorname{Censored}(\\operatorname{Normal}(\\text{mu},~\\text{sigma}),~-2,~2)",
         "\\end{array}",
         "$$",
@@ -316,3 +317,22 @@ def test_custom_dist_repr():
 
     str_repr = model.str_repr(include_params=False)
     assert str_repr == "\n".join(["x ~ CustomDistNormal", "y ~ CustomRandomNormal"])
+
+
+class TestLatexRepr:
+    @staticmethod
+    def simple_model() -> Model:
+        with Model() as simple_model:
+            error = HalfNormal("error", 0.5)
+            alpha_a = Normal("alpha_a", 0, 1)
+            Normal("y", alpha_a, error)
+        return simple_model
+
+    def test_latex_escaped_underscore(self):
+        """
+        Ensures that all underscores in model variable names are properly escaped for LaTeX representation
+        """
+        model = self.simple_model()
+        model_str = model.str_repr(formatting="latex")
+        assert "\\_" in model_str
+        assert "_" not in model_str.replace("\\_", "")