Allow draws from Weibull, MvStudentT, LKJCorr and LKJCholeskyCovRV in alternative backends

.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' &&

conda-envs/environment-jax.yml → conda-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

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):

pymc/distributions/distribution.py

@@ -383,6 +383,14 @@ def batch_ndim(self, node: Apply) -> int:
         out_ndim = max(getattr(out.type, "ndim", 0) for out in node.outputs)
         return out_ndim - self.ndim_supp
 
+    def rebuild_rv(self, *args, **kwargs):
+        """Rebuild the RandomVariable with new inputs."""
+        if not hasattr(self, "rv_op"):
+            raise NotImplementedError(
+                f"SymbolicRandomVariable {self} without `rv_op` method cannot be rebuilt automatically."
+            )
+        return self.rv_op(*args, **kwargs)
+
 
 @_change_dist_size.register(SymbolicRandomVariable)
 def change_symbolic_rv_size(op: SymbolicRandomVariable, rv, new_size, expand) -> TensorVariable:
@@ -400,10 +408,10 @@ def change_symbolic_rv_size(op: SymbolicRandomVariable, rv, new_size, expand) ->
 
     params = op.dist_params(rv.owner)
 
-    if expand:
+    if expand and not rv_size_is_none(size):
         new_size = tuple(new_size) + tuple(size)
 
-    return op.rv_op(*params, size=new_size)
+    return op.rebuild_rv(*params, size=new_size)
 
 
 class Distribution(metaclass=DistributionMeta):