Make more distributions symbolic so they work in different backends

pymc-devs · Feb 27, 2025 · 02a5a36 · 02a5a36
1 parent 9f52e3d
commit 02a5a36
Show file tree

Hide file tree

Showing 8 changed files with 214 additions and 187 deletions.
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -281,7 +281,7 @@ jobs:
           name: ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false
 
-  external_samplers:
+  alternative_backends:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
@@ -290,7 +290,11 @@ jobs:
         floatx: [float64]
         python-version: ["3.13"]
         test-subset:
-          - tests/sampling/test_jax.py tests/sampling/test_mcmc_external.py
+          - |
+            tests/distributions/test_random_alternative_backends.py
+            tests/sampling/test_jax.py
+            tests/sampling/test_mcmc_external.py
+
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
@@ -305,7 +309,7 @@ jobs:
           persist-credentials: false
       - uses: mamba-org/setup-micromamba@v2
         with:
-          environment-file: conda-envs/environment-jax.yml
+          environment-file: conda-envs/environment-alternative-backends.yml
           create-args: >-
             python=${{matrix.python-version}}
           environment-name: pymc-test
@@ -324,7 +328,7 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
           env_vars: TEST_SUBSET
-          name: JAX tests - ${{ matrix.os }} ${{ matrix.floatx }}
+          name: Alternative backend tests - ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false
 
   float32:
@@ -378,7 +382,7 @@ jobs:
   all_tests:
     if: ${{ always() }}
     runs-on: ubuntu-latest
-    needs: [ changes, ubuntu, windows, macos, external_samplers, float32 ]
+    needs: [ changes, ubuntu, windows, macos, alternative_backends, float32 ]
     steps:
       - name: Check build matrix status
         if: ${{ needs.changes.outputs.changes == 'true' &&

diff --git a/conda-envs/environment-jax.yml → ...envs/environment-alternative-backends.yml b/conda-envs/environment-jax.yml → ...envs/environment-alternative-backends.yml
@@ -10,6 +10,8 @@ dependencies:
 - cachetools>=4.2.1
 - cloudpickle
 - zarr>=2.5.0,<3
+- numba
+- nutpie >= 0.13.4
 # Jaxlib version must not be greater than jax version!
 - blackjax>=1.2.2
 - jax>=0.4.28

diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
@@ -2595,23 +2595,27 @@ def dist(cls, nu, **kwargs):
         return Gamma.dist(alpha=nu / 2, beta=1 / 2, **kwargs)
 
 
-class WeibullBetaRV(RandomVariable):
+class WeibullBetaRV(SymbolicRandomVariable):
     name = "weibull"
-    signature = "(),()->()"
-    dtype = "floatX"
+    extended_signature = "[rng],[size],(),()->[rng],()"
     _print_name = ("Weibull", "\\operatorname{Weibull}")
 
-    def __call__(self, alpha, beta, size=None, **kwargs):
-        return super().__call__(alpha, beta, size=size, **kwargs)
-
     @classmethod
-    def rng_fn(cls, rng, alpha, beta, size) -> np.ndarray:
-        if size is None:
-            size = np.broadcast_shapes(alpha.shape, beta.shape)
-        return np.asarray(beta * rng.weibull(alpha, size=size))
+    def rv_op(cls, alpha, beta, *, rng=None, size=None) -> np.ndarray:
+        alpha = pt.as_tensor(alpha)
+        beta = pt.as_tensor(beta)
+        rng = normalize_rng_param(rng)
+        size = normalize_size_param(size)
 
+        if rv_size_is_none(size):
+            size = implicit_size_from_params(alpha, beta, ndims_params=cls.ndims_params)
 
-weibull_beta = WeibullBetaRV()
+        next_rng, raw_weibull = pt.random.weibull(alpha, size=size, rng=rng).owner.outputs
+        draws = beta * raw_weibull
+        return cls(
+            inputs=[rng, size, alpha, beta],
+            outputs=[next_rng, draws],
+        )(rng, size, alpha, beta)
 
 
 class Weibull(PositiveContinuous):
@@ -2660,7 +2664,8 @@ class Weibull(PositiveContinuous):
         Scale parameter (beta > 0).
     """
 
-    rv_op = weibull_beta
+    rv_type = WeibullBetaRV
+    rv_op = WeibullBetaRV.rv_op
 
     @classmethod
     def dist(cls, alpha, beta, *args, **kwargs):