Fix hang in VonMises rejection sampling for small values of concentration

[julian-urban]

Fixes #88443

Forces the internal dtype of torch.distributions.von_mises.VonMises to be torch.double and mirrors the numpy implementation of the second order Taylor expansion for concentration < 1e-5. Samples and log probs are returned with dtype of argument loc.

In principle one could also use masking in the rejection sampler to return uniformly distributed numbers for concentration < 1e-8, as in numpy. This may be slightly more efficient, but isn't required to solve the hanging issue.

cc @fritzo @neerajprad @alicanb @nikitaved

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/114498

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8d068d3 with merge base b27565a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[julian-urban]

@pytorchbot label "topic: not user facing"

[fritzo]

Thanks for fixing this! Could we move your logic to three @lazy_propertys, to keep the other code paths entirely single precision?

[julian-urban]

@fritzo Thanks for picking it up, please check my latest commit. I followed your suggestion of implementing the _loc and _concentration attributes with double prec as lazy properties and similarly _proposal_r, also replacing the if/else with torch.where to make sure it works with tensor-valued loc and concentration. The hard-coded double prec stuff should now be confined to calls of sample() only.

I'm not sure if I was supposed to resolve the above conversation already, this is my first PR and I'm not familiar with the proper etiquette. Apologies.

[fritzo]

Looks great @julian-urban! It might be worth adding a line comment in .sample() explaining why sampling is performed using double precision, but feel free to leave as is.

[fritzo]

Looks like a legit lint error

  Warning (UFMT) format
    Run `lintrunner -a` to apply this patch.

    You can run `lintrunner -a` to apply this patch.

    167  167 |         """
    168  168 |         shape = self._extended_shape(sample_shape)
    169  169 |         x = torch.empty(shape, dtype=self._loc.dtype, device=self.loc.device)
    169      |-        return _rejection_sample(self._loc, self._concentration, self._proposal_r, x).to(self.loc.dtype)
         170 |+        return _rejection_sample(
         171 |+            self._loc, self._concentration, self._proposal_r, x
         172 |+        ).to(self.loc.dtype)
    171  173 | 
    172  174 |     def expand(self, batch_shape):
    173  175 |         try:

[julian-urban]

@fritzo whoopsie. The test passes now and I added a clarifying comment to sample().

[fritzo]

@ezyang could you please merge this?

[fritzo]

@pytorchmergebot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here