Commits on Mar 25, 2022

1. Adding experimental nnc_compile option to NUTS and HMC (#1385) …

   Summary:
   ### Motivation
   With the first Beta release of [functorch](https://github.com/pytorch/functorch), we can begin to merge in our BM-NNC integration prototype, which uses NNC to JIT compile part of the algorithm to accelerate inferences.
   
   ### Changes proposed
   - `functorch>=0.1.0` is added to out list of dependencies
   - Because NNC is yet to support control flow primitives, in NUTS, NNC is applied on the base case of recursive tree building algorithm. In HMC, NNC is applied on a single leapfrog step.
   - We use `torch.Tensor` instead of raw scalars for some variables because TorchScript tracer requires inputs/outputs to be of the same type (?) (i.e., we can't return a tuple of a mixture of `Tensor`s and `float`s)
   - All of the NNC utils are put into `beanmachine.ppl.experimental.nnc.utils`, which will throw a warning when it's being imported for the first time.
   - The docstring of HMC & NUTS classes are updated as well
   
   To try NNC out, simply set `nnc_compile` to `True` when initializing the inference class, e.g.
   
   ```
   nuts = bm.GlobalNoUTurnSampler(nnc_compile=True)
   nuts.infer(...)  # same arguments as usual
   ```
   
   Pull Request resolved: #1385
   
   Test Plan:
   I've added a simple sanity check to cover the NNC compile option on NUTS and HMC, which can be run with
   ```
   buck test //beanmachine/beanmachine/ppl:test-ppl -- nnc
   ```
   or equivalently, for OSS:
   ```
   pytest src/beanmachine/ppl/experimental/tests/nnc_test.py
   ```
   
   Reviewed By: jpchen
   
   Differential Revision: D35127777
   
   Pulled By: horizon-blue
   
   fbshipit-source-id: 578c62b7ef10555fc0f58849e47a6a22c3e51f68

   

   horizon-blue authored and facebook-github-bot committed Mar 25, 2022
   Configuration menu

   • View commit details
   • Copy full SHA for cf0f890
   • Browse repository at this point

   Copy the full SHA

   cf0f890 View commit details

   Browse the repository at this point in the history