-
-
Notifications
You must be signed in to change notification settings - Fork 2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add Truncated normal dispatches #7506
base: main
Are you sure you want to change the base?
Conversation
[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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This fails with a NotImplementedError: No JAX implementation for the given distribution: truncated_normal
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The dispatch file needs to be imported when pymc is imported in order to be registered
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #7506 +/- ##
==========================================
+ Coverage 88.58% 92.43% +3.84%
==========================================
Files 103 103
Lines 17104 17109 +5
==========================================
+ Hits 15152 15814 +662
+ Misses 1952 1295 -657 |
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can't pass sigma or mu as parameters in jax.random.truncated_normal
[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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
|
||
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
TruncatedNormal
distribution for forward sampling #7489Checklist
Type of change
📚 Documentation preview 📚: https://pymc--7506.org.readthedocs.build/en/7506/