From 80cc4f250233d2574b3c43fe2b2b35119d354176 Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sat, 29 Nov 2025 21:12:56 +0100
Subject: [PATCH 01/16] ci init

---
 .github/worflows/test.yml | 52 +++++++++++++++++++++++++++++++++
 pyproject.toml            | 60 ++++++++++++++++++---------------------
 2 files changed, 80 insertions(+), 32 deletions(-)
 create mode 100644 .github/worflows/test.yml

diff --git a/.github/worflows/test.yml b/.github/worflows/test.yml
new file mode 100644
index 0000000..66c41d3
--- /dev/null
+++ b/.github/worflows/test.yml
@@ -0,0 +1,52 @@
+name: Run tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
+      with:
+        python-version: "3.13"
+    - uses: pre-commit/action@v3.0.1
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+
+    name: Set up Python ${{ matrix.python-version }}
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          channels: conda-forge, defaults
+          channel-priority: true
+          python-version: ${{ matrix.python-version }}
+          auto-update-conda: true
+
+      - name: Install distributions
+        shell: bash -l {0}
+        run: |
+          conda install pip
+          pip install -e .[test]
+          python --version
+          conda list
+          pip freeze
+      - name: Run tests
+        shell: bash -l {0}
+        run: |
+          python -m pytest -vv --cov=distributions --cov-report=term --cov-report=xml tests
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
diff --git a/pyproject.toml b/pyproject.toml
index 87ae294..ceafeda 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,9 +6,7 @@ build-backend = "flit_core.buildapi"
 name = "distributions"
 readme = "README.md"
 requires-python = ">=3.11"
-authors = [
-    {name = "pymc-devs", email = "pymc.devs@gmail.com"}
-]
+authors = [{ name = "pymc-devs", email = "pymc.devs@gmail.com" }]
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Science/Research",
@@ -23,10 +21,11 @@ classifiers = [
 ]
 dynamic = ["version"]
 description = "PyTensor powered distributions."
-dependencies = [
-  "numpy>=2.0",
-  "pytensor>=2.32.0",
-]
+dependencies = ["numpy>=2.0", "pytensor>=2.32.0"]
+
+
+[project.optional-dependencies]
+test = ["pytest", "pytest-cov"]
 
 [tool.flit.module]
 name = "distributions"
@@ -38,22 +37,19 @@ documentation = "https://distributions.readthedocs.io"
 funding = "https://opencollective.com/pymc"
 
 
-[tool.black]
-line-length = 100
-
-[tool.isort]
-profile = "black"
-include_trailing_comma = true
-use_parentheses = true
-multi_line_output = 3
-line_length = 100
-
 [tool.pydocstyle]
 convention = "numpy"
 
 [tool.pytest.ini_options]
-testpaths = [
-    "tests",
+testpaths = ["tests"]
+addopts = [
+    "-v",
+    "--strict-markers",
+    "--strict-config",
+    "--color=yes",
+    "--cov=src",
+    "--cov=tests",
+    "--cov-report=term-missing",
 ]
 
 [tool.ruff]
@@ -61,23 +57,23 @@ line-length = 100
 
 [tool.ruff.lint]
 select = [
-    "F",  # Pyflakes
-    "E",  # Pycodestyle
-    "W",  # Pycodestyle
-    "D",  # pydocstyle
+    "F",   # Pyflakes
+    "E",   # Pycodestyle
+    "W",   # Pycodestyle
+    "D",   # pydocstyle
     "NPY", # numpy specific rules
     "UP",  # pyupgrade
-    "I",  # isort
+    "I",   # isort
     "PL",  # Pylint
-    "TID",  # Absolute imports
+    "TID", # Absolute imports
 ]
 ignore = [
-    "PLR0912",  # too many branches
-    "PLR0913",  # too many arguments
-    "PLR2004",  # magic value comparison
-    "PLR0915",  # too many statements
-    "PLC0415",  # import outside of top level
-    "D1"    # Missing docstring
+    "PLR0912", # too many branches
+    "PLR0913", # too many arguments
+    "PLR2004", # magic value comparison
+    "PLR0915", # too many statements
+    "PLC0415", # import outside of top level
+    "D1",      # Missing docstring
 ]
 
 [tool.ruff.lint.per-file-ignores]
@@ -90,7 +86,7 @@ ignore = [
 convention = "numpy"
 
 [tool.ruff.lint.flake8-tidy-imports]
-ban-relative-imports = "all"  # Disallow all relative imports.
+ban-relative-imports = "all" # Disallow all relative imports.
 
 [tool.ruff.format]
 docstring-code-format = false

From 8c77e216b64fd8c30899891ccc5fa33ae55e691b Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sat, 29 Nov 2025 21:15:05 +0100
Subject: [PATCH 02/16] rename

---
 .github/{worflows => workflows}/test.yml | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename .github/{worflows => workflows}/test.yml (100%)

diff --git a/.github/worflows/test.yml b/.github/workflows/test.yml
similarity index 100%
rename from .github/worflows/test.yml
rename to .github/workflows/test.yml

From 8c82e8a31ec7eda9f9dbf751880fa957466c4c6a Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sat, 29 Nov 2025 21:18:01 +0100
Subject: [PATCH 03/16] test with one version for now

---
 .github/workflows/test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 66c41d3..e1cc040 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.11", "3.12", "3.13"]
+        python-version: ["3.13"]
 
     name: Set up Python ${{ matrix.python-version }}
     steps:

From 6dd8ee0c9e7da6aff9d648fb2beb01471a3d8de6 Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sat, 29 Nov 2025 21:26:28 +0100
Subject: [PATCH 04/16] fix src

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index ceafeda..563e8a0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -47,7 +47,7 @@ addopts = [
     "--strict-markers",
     "--strict-config",
     "--color=yes",
-    "--cov=src",
+    "--cov=distributions",
     "--cov=tests",
     "--cov-report=term-missing",
 ]

From 89581c0715607e49d9db717f52040401d0453ed9 Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sat, 29 Nov 2025 21:55:40 +0100
Subject: [PATCH 05/16] try to fix tests

---
 .gitignore                    |  5 ++-
 distributions/betabinomial.py | 32 ++++++++++++++++---
 distributions/optimization.py | 60 ++++++++++++++++++++++++++++++-----
 tests/test_wald.py            |  4 +++
 4 files changed, 88 insertions(+), 13 deletions(-)

diff --git a/.gitignore b/.gitignore
index ed8ebf5..09e5781 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,4 @@
-__pycache__
\ No newline at end of file
+__pycache__
+
+.coverage
+coverage.xml
\ No newline at end of file
diff --git a/distributions/betabinomial.py b/distributions/betabinomial.py
index 8689388..e630fc7 100644
--- a/distributions/betabinomial.py
+++ b/distributions/betabinomial.py
@@ -9,11 +9,35 @@ def mean(n, alpha, beta):
 
 
 def mode(n, alpha, beta):
-    return pt.clip(
-        pt.floor((n + 1) * ((alpha - 1) / (alpha + beta - 2))),
-        0,
-        n,
+    # The standard formula only applies when alpha > 1 and beta > 1
+    # For other cases, the mode is at the boundaries or distribution is uniform
+    standard_mode = pt.floor((n + 1) * ((alpha - 1) / (alpha + beta - 2)))
+
+    # Handle different cases:
+    # - alpha < 1 and beta < 1: U-shaped, modes at 0 and n (return 0)
+    # - alpha <= 1 and beta > 1: mode at 0
+    # - alpha > 1 and beta <= 1: mode at n
+    # - alpha = beta = 1: uniform, any value valid (return floor(n/2))
+    # - alpha > 1 and beta > 1: use standard formula
+
+    result = pt.switch(
+        pt.and_(pt.gt(alpha, 1), pt.gt(beta, 1)),
+        # Standard case: alpha > 1 and beta > 1
+        pt.clip(standard_mode, 0, n),
+        pt.switch(
+            pt.and_(pt.le(alpha, 1), pt.le(beta, 1)),
+            # Both <= 1: U-shaped or uniform, mode at boundary (use 0)
+            0,
+            pt.switch(
+                pt.le(alpha, 1),
+                # alpha <= 1, beta > 1: mode at 0
+                0,
+                # alpha > 1, beta <= 1: mode at n
+                n,
+            ),
+        ),
     )
+    return result
 
 
 def median(n, alpha, beta):
diff --git a/distributions/optimization.py b/distributions/optimization.py
index 3fc70f7..567e2f4 100644
--- a/distributions/optimization.py
+++ b/distributions/optimization.py
@@ -1,3 +1,5 @@
+import math
+
 import pytensor.tensor as pt
 
 from distributions.helper import ppf_bounds_cont, ppf_bounds_disc
@@ -48,15 +50,57 @@ def func(x):
     return ppf_bounds_cont(0.5 * (left + right), q, lower, upper)
 
 
+def _is_unbounded(upper):
+    """Check if upper bound is infinite at Python level."""
+    # Check for Python float infinity
+    if isinstance(upper, float) and math.isinf(upper):
+        return True
+    # Check for numpy infinity
+    try:
+        import numpy as np
+
+        if isinstance(upper, (np.floating, np.integer)) and np.isinf(upper):
+            return True
+    except (ImportError, TypeError):
+        pass
+    return False
+
+
 def find_ppf_discrete(q, lower, upper, cdf, *params):
     """
-    Compute the inverse CDF using the bisection method.
+    Compute the inverse CDF for discrete distributions.
 
-    The continuous bisection method finds where CDF(x) ≈ q. For discrete distributions,
-    we round to the nearest integer and then check if we need to adjust.
+    For bounded support, uses direct search over all values (fast).
+    For unbounded support, uses bisection method (works with infinite bounds).
     """
-    rounded_k = pt.round(find_ppf(q, lower, upper, cdf, *params))
-    # return ppf_bounds_disc(rounded_k, q, lower, upper)
-    cdf_k = cdf(rounded_k, *params)
-    rounded_k = pt.switch(pt.lt(cdf_k, q), rounded_k + 1, rounded_k)
-    return ppf_bounds_disc(rounded_k, q, lower, upper)
+    if _is_unbounded(upper):
+        # Unbounded case: use bisection method
+        rounded_k = pt.round(find_ppf(q, lower, upper, cdf, *params))
+        cdf_k = cdf(rounded_k, *params)
+        rounded_k = pt.switch(pt.lt(cdf_k, q), rounded_k + 1, rounded_k)
+        return ppf_bounds_disc(rounded_k, q, lower, upper)
+
+    # Bounded case: direct search over all values
+    q = pt.as_tensor_variable(q)
+
+    # Create array of all possible values in support
+    k_vals = pt.arange(lower, upper + 1)
+
+    # Compute CDF for all values - shape: (n_support,)
+    cdf_vals = cdf(k_vals, *params)
+
+    # Use a small tolerance for floating point comparison
+    eps = 1e-10
+
+    if q.ndim == 0:
+        # Scalar case
+        exceeds_q = pt.ge(cdf_vals, q - eps)
+        first_idx = pt.argmax(exceeds_q)
+        result = k_vals[first_idx]
+    else:
+        # Array case - need broadcasting
+        exceeds_q = pt.ge(cdf_vals[:, None], q[None, :] - eps)
+        first_idx = pt.argmax(exceeds_q, axis=0)
+        result = k_vals[first_idx]
+
+    return ppf_bounds_disc(result, q, lower, upper)
diff --git a/tests/test_wald.py b/tests/test_wald.py
index 9d3e2ed..e3bbcc5 100644
--- a/tests/test_wald.py
+++ b/tests/test_wald.py
@@ -29,4 +29,8 @@ def test_wald_vs_scipy(params, sp_params):
         sp_params=sp_params,
         support=support,
         name="wald",
+        # Slightly higher tolerance for SF/logCDF due to numerical precision
+        # when CDF is very close to 1 (error is ~1e-9 absolute, just over 1e-6 relative)
+        sf_rtol=1.1e-6,
+        logcdf_rtol=1.1e-6,
     )

From 3a516ed30577867426aed0a4184cfb232dd6c913 Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sat, 29 Nov 2025 21:56:32 +0100
Subject: [PATCH 06/16] ruff-check

---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 021bba5..d55802f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,7 +4,7 @@ repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.14.1
     hooks:
-      - id: ruff
+      - id: ruff-check
         args: [ --fix, --exit-non-zero-on-fix ]
       - id: ruff-format
         types_or: [ python, pyi ]

From 393a6fa145b749551675fa107067a4e930edae8b Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sun, 30 Nov 2025 11:58:19 +0100
Subject: [PATCH 07/16] use pytensor

---
 distributions/optimization.py | 23 +++++++++--------------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/distributions/optimization.py b/distributions/optimization.py
index 567e2f4..15a534b 100644
--- a/distributions/optimization.py
+++ b/distributions/optimization.py
@@ -1,5 +1,3 @@
-import math
-
 import pytensor.tensor as pt
 
 from distributions.helper import ppf_bounds_cont, ppf_bounds_disc
@@ -51,19 +49,16 @@ def func(x):
 
 
 def _is_unbounded(upper):
-    """Check if upper bound is infinite at Python level."""
-    # Check for Python float infinity
-    if isinstance(upper, float) and math.isinf(upper):
-        return True
-    # Check for numpy infinity
-    try:
-        import numpy as np
+    """Check if upper bound is infinite using pytensor.
 
-        if isinstance(upper, (np.floating, np.integer)) and np.isinf(upper):
-            return True
-    except (ImportError, TypeError):
-        pass
-    return False
+    Evaluates pt.isinf at graph-build time for constants.
+    For symbolic expressions, assumes unbounded (safe default for bisection).
+    """
+    try:
+        upper_t = pt.as_tensor_variable(upper)
+        return bool(pt.isinf(upper_t).eval())
+    except Exception:
+        return True
 
 
 def find_ppf_discrete(q, lower, upper, cdf, *params):

From 4e11dc97f12185709e689df866757edf237ee39b Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sun, 30 Nov 2025 12:04:16 +0100
Subject: [PATCH 08/16] mode can be nan

---
 distributions/betabinomial.py | 26 +++++---------------------
 tests/test_betabinomial.py    | 19 ++++++++++++-------
 2 files changed, 17 insertions(+), 28 deletions(-)

diff --git a/distributions/betabinomial.py b/distributions/betabinomial.py
index e630fc7..c8caea5 100644
--- a/distributions/betabinomial.py
+++ b/distributions/betabinomial.py
@@ -10,32 +10,16 @@ def mean(n, alpha, beta):
 
 def mode(n, alpha, beta):
     # The standard formula only applies when alpha > 1 and beta > 1
-    # For other cases, the mode is at the boundaries or distribution is uniform
+    # For other cases, the mode is not unique (U-shaped, monotonic, or uniform)
+    # Following the convention used in this package, return NaN when mode is not unique
     standard_mode = pt.floor((n + 1) * ((alpha - 1) / (alpha + beta - 2)))
 
-    # Handle different cases:
-    # - alpha < 1 and beta < 1: U-shaped, modes at 0 and n (return 0)
-    # - alpha <= 1 and beta > 1: mode at 0
-    # - alpha > 1 and beta <= 1: mode at n
-    # - alpha = beta = 1: uniform, any value valid (return floor(n/2))
-    # - alpha > 1 and beta > 1: use standard formula
-
     result = pt.switch(
         pt.and_(pt.gt(alpha, 1), pt.gt(beta, 1)),
-        # Standard case: alpha > 1 and beta > 1
+        # Standard case: alpha > 1 and beta > 1 - unique mode exists
         pt.clip(standard_mode, 0, n),
-        pt.switch(
-            pt.and_(pt.le(alpha, 1), pt.le(beta, 1)),
-            # Both <= 1: U-shaped or uniform, mode at boundary (use 0)
-            0,
-            pt.switch(
-                pt.le(alpha, 1),
-                # alpha <= 1, beta > 1: mode at 0
-                0,
-                # alpha > 1, beta <= 1: mode at n
-                n,
-            ),
-        ),
+        # Otherwise: mode is not unique, return NaN
+        pt.nan,
     )
     return result
 
diff --git a/tests/test_betabinomial.py b/tests/test_betabinomial.py
index abd9824..ad2144a 100644
--- a/tests/test_betabinomial.py
+++ b/tests/test_betabinomial.py
@@ -9,16 +9,20 @@
 
 
 @pytest.mark.parametrize(
-    "params, sp_params",
+    "params, sp_params, skip_mode",
     [
-        ([10, 2.0, 3.0], {"n": 10, "a": 2.0, "b": 3.0}),
-        ([5, 1.0, 1.0], {"n": 5, "a": 1.0, "b": 1.0}),
-        ([20, 0.5, 0.5], {"n": 20, "a": 0.5, "b": 0.5}),
-        ([15, 5.0, 2.0], {"n": 15, "a": 5.0, "b": 2.0}),
-        ([100, 20.0, 20.0], {"n": 100, "a": 20.0, "b": 20.0}),
+        # alpha > 1 and beta > 1: unique mode exists
+        ([10, 2.0, 3.0], {"n": 10, "a": 2.0, "b": 3.0}, False),
+        # alpha = beta = 1: uniform, mode not unique
+        ([5, 1.0, 1.0], {"n": 5, "a": 1.0, "b": 1.0}, True),
+        # alpha < 1 and beta < 1: U-shaped, mode not unique
+        ([20, 0.5, 0.5], {"n": 20, "a": 0.5, "b": 0.5}, True),
+        # alpha > 1 and beta > 1: unique mode exists
+        ([15, 5.0, 2.0], {"n": 15, "a": 5.0, "b": 2.0}, False),
+        ([100, 20.0, 20.0], {"n": 100, "a": 20.0, "b": 20.0}, False),
     ],
 )
-def test_betabinomial_vs_scipy(params, sp_params):
+def test_betabinomial_vs_scipy(params, sp_params, skip_mode):
     """Test BetaBinomial distribution against scipy."""
     n_param = pt.constant(params[0], dtype="int64")
     alpha_param = pt.constant(params[1], dtype="float64")
@@ -34,4 +38,5 @@ def test_betabinomial_vs_scipy(params, sp_params):
         support=support,
         is_discrete=True,
         name="betabinomial",
+        skip_mode=skip_mode,
     )

From b5352bb47869745c6b4247ffe23b12ef026c3d59 Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sun, 30 Nov 2025 12:11:12 +0100
Subject: [PATCH 09/16] improve optimization

---
 distributions/optimization.py | 39 ++++++++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/distributions/optimization.py b/distributions/optimization.py
index 15a534b..d535ee9 100644
--- a/distributions/optimization.py
+++ b/distributions/optimization.py
@@ -1,3 +1,5 @@
+import math
+
 import pytensor.tensor as pt
 
 from distributions.helper import ppf_bounds_cont, ppf_bounds_disc
@@ -48,16 +50,33 @@ def func(x):
     return ppf_bounds_cont(0.5 * (left + right), q, lower, upper)
 
 
-def _is_unbounded(upper):
-    """Check if upper bound is infinite using pytensor.
+def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
+    """Check if bisection should be used instead of direct search.
+
+    Uses bisection if:
+    - Upper bound is infinite
+    - Range is too wide (> max_direct_search_size)
+    - Values cannot be evaluated at graph-build time
 
-    Evaluates pt.isinf at graph-build time for constants.
-    For symbolic expressions, assumes unbounded (safe default for bisection).
+    Evaluates bounds at graph-build time for constants.
     """
     try:
+        lower_t = pt.as_tensor_variable(lower)
         upper_t = pt.as_tensor_variable(upper)
-        return bool(pt.isinf(upper_t).eval())
+        lower_val = float(lower_t.eval())
+        upper_val = float(upper_t.eval())
+
+        # Check for infinite bounds
+        if not (math.isfinite(lower_val) and math.isfinite(upper_val)):
+            return True
+
+        # Check if range is too wide
+        if (upper_val - lower_val) > max_direct_search_size:
+            return True
+
+        return False
     except Exception:
+        # If evaluation fails, use bisection as safe default
         return True
 
 
@@ -65,17 +84,17 @@ def find_ppf_discrete(q, lower, upper, cdf, *params):
     """
     Compute the inverse CDF for discrete distributions.
 
-    For bounded support, uses direct search over all values (fast).
-    For unbounded support, uses bisection method (works with infinite bounds).
+    For narrow bounded support, uses direct search over all values (fast).
+    For unbounded or wide support, uses bisection method.
     """
-    if _is_unbounded(upper):
-        # Unbounded case: use bisection method
+    if _should_use_bisection(lower, upper):
+        # Use bisection method for unbounded or wide ranges
         rounded_k = pt.round(find_ppf(q, lower, upper, cdf, *params))
         cdf_k = cdf(rounded_k, *params)
         rounded_k = pt.switch(pt.lt(cdf_k, q), rounded_k + 1, rounded_k)
         return ppf_bounds_disc(rounded_k, q, lower, upper)
 
-    # Bounded case: direct search over all values
+    # Bounded case with narrow range: direct search over all values
     q = pt.as_tensor_variable(q)
 
     # Create array of all possible values in support

From c78a16c2004f8d5ebd0c89d7381c13090426ecee Mon Sep 17 00:00:00 2001
From: juanitorduz <juanitorduz@gmail.com>
Date: Sun, 30 Nov 2025 21:05:50 +0100
Subject: [PATCH 10/16] cleanup

---
 distributions/optimization.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/distributions/optimization.py b/distributions/optimization.py
index d535ee9..94c11bf 100644
--- a/distributions/optimization.py
+++ b/distributions/optimization.py
@@ -50,21 +50,29 @@ def func(x):
     return ppf_bounds_cont(0.5 * (left + right), q, lower, upper)
 
 
+def _get_constant_value(x):
+    """Extract numeric value from a constant (pytensor or Python scalar).
+
+    Returns the float value if x is a constant, raises exception otherwise.
+    """
+    # PyTensor constant - access underlying data directly
+    if hasattr(x, "data"):
+        return float(x.data)
+    # Python/numpy scalar - convert directly
+    return float(x)
+
+
 def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
     """Check if bisection should be used instead of direct search.
 
     Uses bisection if:
-    - Upper bound is infinite
+    - Bounds are infinite
     - Range is too wide (> max_direct_search_size)
-    - Values cannot be evaluated at graph-build time
-
-    Evaluates bounds at graph-build time for constants.
+    - Bounds are symbolic (not constants)
     """
     try:
-        lower_t = pt.as_tensor_variable(lower)
-        upper_t = pt.as_tensor_variable(upper)
-        lower_val = float(lower_t.eval())
-        upper_val = float(upper_t.eval())
+        lower_val = _get_constant_value(lower)
+        upper_val = _get_constant_value(upper)
 
         # Check for infinite bounds
         if not (math.isfinite(lower_val) and math.isfinite(upper_val)):
@@ -76,7 +84,7 @@ def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
 
         return False
     except Exception:
-        # If evaluation fails, use bisection as safe default
+        # If extraction fails (symbolic tensor), use bisection as safe default
         return True
 
 

From e361dbf97d7efa52ac9fd249d22f5efe9c83eeab Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Wed, 14 Jan 2026 23:28:12 +0100
Subject: [PATCH 11/16] feedback

---
 .github/workflows/test.yml    |  2 +-
 distributions/optimization.py | 71 +++++++++++++++++++----------------
 2 files changed, 39 insertions(+), 34 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e1cc040..a319fc5 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -24,7 +24,7 @@ jobs:
 
     name: Set up Python ${{ matrix.python-version }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/distributions/optimization.py b/distributions/optimization.py
index 4da8cea..994be6a 100644
--- a/distributions/optimization.py
+++ b/distributions/optimization.py
@@ -50,43 +50,53 @@ def func(x):
     return ppf_bounds_cont(0.5 * (left + right), q, lower, upper)
 
 
-def _get_constant_value(x):
-    """Extract numeric value from a constant (pytensor or Python scalar).
+def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
+    """Compile-time check to select PPF algorithm for discrete distributions.
 
-    Returns the float value if x is a constant, raises exception otherwise.
-    """
-    # PyTensor constant - access underlying data directly
-    if hasattr(x, "data"):
-        return float(x.data)
-    # Python/numpy scalar - convert directly
-    return float(x)
+    This function inspects bounds at graph-construction time to choose between:
+    - Direct search: Fast for narrow bounded support (e.g., BetaBinomial, Binomial)
+    - Bisection: Required for unbounded or wide support (e.g., Poisson, NegativeBinomial)
 
+    The check happens at Python level during graph construction, not during
+    PyTensor execution. This is intentional: a fully symbolic approach using
+    pt.switch would evaluate both branches, causing performance issues.
 
-def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
-    """Check if bisection should be used instead of direct search.
+    Parameters
+    ----------
+    lower : int, float, or PyTensor constant
+        Lower bound of the distribution support
+    upper : int, float, or PyTensor constant
+        Upper bound of the distribution support
+    max_direct_search_size : int, default 10_000
+        Maximum range size for direct search. Larger ranges use bisection.
 
-    Uses bisection if:
-    - Bounds are infinite
-    - Range is too wide (> max_direct_search_size)
-    - Bounds are symbolic (not constants)
+    Returns
+    -------
+    bool
+        True if bisection should be used, False for direct search.
     """
     try:
-        lower_val = _get_constant_value(lower)
-        upper_val = _get_constant_value(upper)
-
-        # Check for infinite bounds
-        if not (math.isfinite(lower_val) and math.isfinite(upper_val)):
-            return True
-
-        # Check if range is too wide
-        if (upper_val - lower_val) > max_direct_search_size:
-            return True
+        # Extract constant values at graph-build time
+        if hasattr(lower, "data"):
+            lower_val = float(lower.data)
+        else:
+            lower_val = float(lower)
+
+        if hasattr(upper, "data"):
+            upper_val = float(upper.data)
+        else:
+            upper_val = float(upper)
+    except (TypeError, ValueError):
+        # Symbolic (non-constant) bounds - use bisection as safe default
+        return True
 
-        return False
-    except Exception:
-        # If extraction fails (symbolic tensor), use bisection as safe default
+    # Check for infinite bounds
+    if not (math.isfinite(lower_val) and math.isfinite(upper_val)):
         return True
 
+    # Check if range exceeds threshold
+    return (upper_val - lower_val) > max_direct_search_size
+
 
 def find_ppf_discrete(q, lower, upper, cdf, *params):
     """
@@ -126,11 +136,6 @@ def find_ppf_discrete(q, lower, upper, cdf, *params):
         result = k_vals[first_idx]
 
     return ppf_bounds_disc(result, q, lower, upper)
-    rounded_k = pt.round(find_ppf(q, lower, upper, cdf, *params))
-    # return ppf_bounds_disc(rounded_k, q, lower, upper)
-    cdf_k = cdf(rounded_k, *params)
-    rounded_k = pt.switch(pt.lt(cdf_k, q), rounded_k + 1, rounded_k)
-    return ppf_bounds_disc(rounded_k, q, lower, upper)
 
 
 def von_mises_ppf(q, mu, kappa, cdf_func):

From b5760ce315a066570b4fce020dfd3e68820dc1fc Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Wed, 14 Jan 2026 23:30:25 +0100
Subject: [PATCH 12/16] relax test condition

---
 tests/test_logitnormal.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_logitnormal.py b/tests/test_logitnormal.py
index 7429d53..517203b 100644
--- a/tests/test_logitnormal.py
+++ b/tests/test_logitnormal.py
@@ -11,7 +11,7 @@
     "params",
     [
         [0.0, 1.0],  # Standard logit-normal (centered)
-        [0.0, 0.001],  # Narrower distribution
+        [0.0, 0.5],  # Narrower distribution (sigma=0.001 is too extreme for numerical integration)
         [1.0, 1.0],  # Shifted right (mode > 0.5)
         [-1.0, 1.0],  # Shifted left (mode < 0.5)
         [0.0, 2.0],  # Wider distribution (approaches U-shape)

From 215c52ec1ef1f4ad0d62ce3a68cbda422b337e7a Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Wed, 14 Jan 2026 23:32:15 +0100
Subject: [PATCH 13/16] rename

---
 .github/workflows/{test.yml => ci.yml} | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)
 rename .github/workflows/{test.yml => ci.yml} (98%)

diff --git a/.github/workflows/test.yml b/.github/workflows/ci.yml
similarity index 98%
rename from .github/workflows/test.yml
rename to .github/workflows/ci.yml
index a319fc5..9fad603 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/ci.yml
@@ -1,4 +1,4 @@
-name: Run tests
+name: ci
 
 on:
   push:

From b23a326af23d3d2420dde892b5e272dcedb6ee50 Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Wed, 14 Jan 2026 23:33:42 +0100
Subject: [PATCH 14/16] rename

---
 .github/workflows/ci.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 9fad603..c797f2b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -22,13 +22,13 @@ jobs:
       matrix:
         python-version: ["3.13"]
 
-    name: Set up Python ${{ matrix.python-version }}
+    name: Test ${{ matrix.python-version }}
     steps:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
-      - name: Set up Python ${{ matrix.python-version }}
+      - name: Setup environment ${{ matrix.python-version }}
         uses: conda-incubator/setup-miniconda@v2
         with:
           channels: conda-forge, defaults

From d3fcd9276262d23f0edaeffb1859ca671d7c7626 Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Wed, 14 Jan 2026 23:34:16 +0100
Subject: [PATCH 15/16] rename

---
 .github/workflows/ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index c797f2b..6c0303b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -22,7 +22,7 @@ jobs:
       matrix:
         python-version: ["3.13"]
 
-    name: Test ${{ matrix.python-version }}
+    name: tests ${{ matrix.python-version }}
     steps:
       - uses: actions/checkout@v4
         with:

From 31727204549a556f0c4d86a0ade8229b03d606f4 Mon Sep 17 00:00:00 2001
From: Juan Orduz <juanitorduz@gmail.com>
Date: Wed, 14 Jan 2026 23:48:04 +0100
Subject: [PATCH 16/16] feedback cursor

---
 distributions/betabinomial.py | 32 +++++++++++++++++++++-----------
 distributions/optimization.py | 23 +++++++++++++++++++++--
 tests/test_betabinomial.py    |  4 ++++
 3 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/distributions/betabinomial.py b/distributions/betabinomial.py
index c8caea5..747fd87 100644
--- a/distributions/betabinomial.py
+++ b/distributions/betabinomial.py
@@ -9,18 +9,28 @@ def mean(n, alpha, beta):
 
 
 def mode(n, alpha, beta):
-    # The standard formula only applies when alpha > 1 and beta > 1
-    # For other cases, the mode is not unique (U-shaped, monotonic, or uniform)
-    # Following the convention used in this package, return NaN when mode is not unique
-    standard_mode = pt.floor((n + 1) * ((alpha - 1) / (alpha + beta - 2)))
-
-    result = pt.switch(
-        pt.and_(pt.gt(alpha, 1), pt.gt(beta, 1)),
-        # Standard case: alpha > 1 and beta > 1 - unique mode exists
-        pt.clip(standard_mode, 0, n),
-        # Otherwise: mode is not unique, return NaN
-        pt.nan,
+    # The mode depends on the shape of the distribution:
+    # - alpha > 1, beta > 1: unimodal, standard formula applies
+    # - alpha = 1, beta > 1: monotonically decreasing, mode is 0
+    # - alpha > 1, beta = 1: monotonically increasing, mode is n
+    # - alpha = 1, beta = 1: uniform, no unique mode (return NaN)
+    # - alpha < 1 or beta < 1 (other cases): U-shaped or J-shaped, no unique mode (return NaN)
+    # This follows the same convention as distributions/beta.py
+    n_b, alpha_b, beta_b = pt.broadcast_arrays(n, alpha, beta)
+    result = pt.full_like(alpha_b, pt.nan, dtype="float64")
+
+    # Monotonically decreasing: alpha = 1 and beta > 1 -> mode is 0
+    result = pt.where(pt.eq(alpha_b, 1) & pt.gt(beta_b, 1), 0.0, result)
+    # Monotonically increasing: alpha > 1 and beta = 1 -> mode is n
+    result = pt.where(pt.gt(alpha_b, 1) & pt.eq(beta_b, 1), n_b, result)
+    # Standard unimodal case: alpha > 1 and beta > 1
+    standard_mode = pt.floor((n_b + 1) * ((alpha_b - 1) / (alpha_b + beta_b - 2)))
+    result = pt.where(
+        pt.gt(alpha_b, 1) & pt.gt(beta_b, 1),
+        pt.clip(standard_mode, 0, n_b),
+        result,
     )
+
     return result
 
 
diff --git a/distributions/optimization.py b/distributions/optimization.py
index 994be6a..c5df102 100644
--- a/distributions/optimization.py
+++ b/distributions/optimization.py
@@ -50,7 +50,17 @@ def func(x):
     return ppf_bounds_cont(0.5 * (left + right), q, lower, upper)
 
 
-def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
+def _is_scalar_param(param):
+    """Check if a parameter is a scalar (0-dimensional) at graph-build time."""
+    if hasattr(param, "ndim"):
+        return param.ndim == 0
+    # For Python scalars
+    import numpy as np
+
+    return np.ndim(param) == 0
+
+
+def _should_use_bisection(lower, upper, params, max_direct_search_size=10_000):
     """Compile-time check to select PPF algorithm for discrete distributions.
 
     This function inspects bounds at graph-construction time to choose between:
@@ -67,6 +77,9 @@ def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
         Lower bound of the distribution support
     upper : int, float, or PyTensor constant
         Upper bound of the distribution support
+    params : tuple
+        Distribution parameters - if any are non-scalar, bisection is required
+        to handle broadcasting correctly.
     max_direct_search_size : int, default 10_000
         Maximum range size for direct search. Larger ranges use bisection.
 
@@ -75,6 +88,12 @@ def _should_use_bisection(lower, upper, max_direct_search_size=10_000):
     bool
         True if bisection should be used, False for direct search.
     """
+    # Check if any parameter is non-scalar (array) - direct search doesn't
+    # handle broadcasting correctly, so fall back to bisection
+    for param in params:
+        if not _is_scalar_param(param):
+            return True
+
     try:
         # Extract constant values at graph-build time
         if hasattr(lower, "data"):
@@ -105,7 +124,7 @@ def find_ppf_discrete(q, lower, upper, cdf, *params):
     For narrow bounded support, uses direct search over all values (fast).
     For unbounded or wide support, uses bisection method.
     """
-    if _should_use_bisection(lower, upper):
+    if _should_use_bisection(lower, upper, params):
         # Use bisection method for unbounded or wide ranges
         rounded_k = pt.round(find_ppf(q, lower, upper, cdf, *params))
         cdf_k = cdf(rounded_k, *params)
diff --git a/tests/test_betabinomial.py b/tests/test_betabinomial.py
index ad2144a..2ee5352 100644
--- a/tests/test_betabinomial.py
+++ b/tests/test_betabinomial.py
@@ -20,6 +20,10 @@
         # alpha > 1 and beta > 1: unique mode exists
         ([15, 5.0, 2.0], {"n": 15, "a": 5.0, "b": 2.0}, False),
         ([100, 20.0, 20.0], {"n": 100, "a": 20.0, "b": 20.0}, False),
+        # alpha = 1 and beta > 1: monotonically decreasing, unique mode at 0
+        ([10, 1.0, 3.0], {"n": 10, "a": 1.0, "b": 3.0}, False),
+        # alpha > 1 and beta = 1: monotonically increasing, unique mode at n
+        ([10, 3.0, 1.0], {"n": 10, "a": 3.0, "b": 1.0}, False),
     ],
 )
 def test_betabinomial_vs_scipy(params, sp_params, skip_mode):