Add Truncated normal dispatches #7506
Conversation
tests/dispatch/test_jax.py
Outdated
| [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()) |
There was a problem hiding this comment.
This fails with a NotImplementedError: No JAX implementation for the given distribution: truncated_normal
There was a problem hiding this comment.
The dispatch file needs to be imported when pymc is imported in order to be registered
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #7506 +/- ##
==========================================
+ Coverage 88.58% 92.85% +4.26%
==========================================
Files 103 106 +3
Lines 17104 17612 +508
==========================================
+ Hits 15152 16354 +1202
+ Misses 1952 1258 -694
|
pymc/dispatch/dispatch_jax.py
Outdated
| @jax_funcify.register(TruncatedNormalRV) | ||
| def jax_funcify_TruncatedNormalRV(op, **kwargs): | ||
| def trunc_normal_fn(key, size, mu, sigma, lower, upper): | ||
| return None, jax.random.truncated_normal( |
There was a problem hiding this comment.
Mu and sigma missing and the split rng should be returned, not None.
Check some of the dispatches in PyTensor for a template
There was a problem hiding this comment.
I am using jax.nn.initializers.truncated_normal now, but the tests still fail. Not sure if I have used the rng parameter correctly in tests
pymc/dispatch/dispatch_jax.py
Outdated
| 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) |
There was a problem hiding this comment.
There was a problem hiding this comment.
We can't pass sigma or mu as parameters in jax.random.truncated_normal
There was a problem hiding this comment.
I was able to use jax.random.truncated_normal, had to transform lower and upper
tests/dispatch/test_jax.py
Outdated
| [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()) |
There was a problem hiding this comment.
The two are not expected to match in values, because JAX uses a different implementation than numpy. You can make a TruncatedNormal with a large sigma, and confirm it does not go beyond the bounds as a check
pymc/dispatch/dispatch_jax.py
Outdated
|
|
||
| truncnorm = jax.nn.initializers.truncated_normal(sigma, lower=lower, upper=upper) | ||
|
|
||
| return key, truncnorm(key["jax_state"], size) + mu |
There was a problem hiding this comment.
Adding mu like this is potentially wrong, because when size is None, mu could be larger and we end up with repeated values
There was a problem hiding this comment.
Also you should dispatch on the more specific jax_sample_fn. For the issue with broadcasting, check how we do it here for Normal for example: https://github.com/pymc-devs/pytensor/blob/5d4b0c4b9a1e478dda48e912ee708a9e557e9343/pytensor/link/jax/dispatch/random.py#L147-L173
Description
Add jax dispatch for truncated normal distribution
Related Issue
TruncatedNormaldistribution for forward sampling #7489Checklist
Type of change
📚 Documentation preview 📚: https://pymc--7506.org.readthedocs.build/en/7506/