Add Truncated normal dispatches

.github/workflows/tests.yml

@@ -61,6 +61,7 @@ jobs:
             tests/distributions/test_shape_utils.py
             tests/distributions/test_mixture.py
             tests/test_testing.py
+            tests/dispatch/test_jax.py
 
           - |
             tests/distributions/test_continuous.py

pymc/dispatch/dispatch_jax.py

@@ -0,0 +1,32 @@
+#   Copyright 2024 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.
+import jax
+
+from pytensor.link.jax.dispatch import jax_funcify
+
+from pymc.distributions.continuous import TruncatedNormalRV
+
+
+@jax_funcify.register(TruncatedNormalRV)
+def jax_funcify_TruncatedNormalRV(op, **kwargs):
+    def trunc_normal_fn(key, size, mu, sigma, lower, upper):
+        rng_key = key["jax_state"]
+        rng_key, sampling_key = jax.random.split(rng_key, 2)
+        key["jax_state"] = rng_key
+
+        truncnorm = jax.nn.initializers.truncated_normal(sigma, lower=lower, upper=upper)
+
+        return key, truncnorm(key["jax_state"], size) + mu
+
+    return trunc_normal_fn

tests/dispatch/test_jax.py

@@ -0,0 +1,38 @@
+#   Copyright 2024 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.
+import numpy as np
+import pytest
+
+from pytensor import function
+
+import pymc as pm
+
+from pymc.dispatch import dispatch_jax  # noqa: F401
+
+jax = pytest.importorskip("jax")
+
+
+def test_jax_TruncatedNormal():
+    with pm.Model() as m:
+        f_jax = function(
+            [],
+            [pm.TruncatedNormal("a", 0, 1, lower=-1, upper=2, rng=np.random.default_rng(seed=123))],
+            mode="JAX",
+        )
+        f_py = function(
+            [],
+            [pm.TruncatedNormal("b", 0, 1, lower=-1, upper=2, rng=np.random.default_rng(seed=123))],
+        )
+
+    assert jax.numpy.array_equal(a1=f_py(), a2=f_jax())