Update static docs to describe NNC

docs/framework_topics/mcmc_inference/hamiltonian_monte_carlo.md

@@ -137,9 +137,10 @@ The global variant for adaptive HMC, `GlobalHamiltonianMonteCarlo`, which propos
 bm.GlobalHamiltonianMonteCarlo(
     1.0,  # trajectory length
     initial_step_size=0.1,
-    adapt_step_size=True
-    adapt_mass_matrix=True
-    target_accept_prob=0.8
+    adapt_step_size=True,
+    adapt_mass_matrix=True,
+    target_accept_prob=0.8,
+    nnc_compile=True,
 ).infer(
     queries,
     observations,
@@ -149,6 +150,14 @@ bm.GlobalHamiltonianMonteCarlo(
 )
 ```
 
+:::caution
+
+Functorch's [ahead of time (AOT) autograd compiler](https://pytorch.org/functorch/stable/aot_autograd.html) is used
+by default. If working with a non-static model or unexpected errors are encountered, you may need to manually
+disable the `nnc_compile` flag.
+
+:::
+
 These arguments allow us to decide if we want to tune the step size `adjust_step_size` or covariance matrix `adapt_mass_matrix`. The `target_accept_prob` argument indicates the acceptance probability which should be targeted by the step size tuning algorithm. While the optimal value is 65.1%, higher values have been show to be more robust. As a result, Bean Machine targets an acceptance rate of 0.8 by default.
 
 The parameters to `infer` are described below:

docs/framework_topics/mcmc_inference/no_u_turn_sampler.md

@@ -59,7 +59,6 @@ bm.GlobalNoUTurnSampler(
     adapt_mass_matrix=True,
     multinomial_sampling=True,
     target_accept_prob=0.8,
-    nnc_compile=False,
 ).infer(
     queries,
     observations,
@@ -69,6 +68,14 @@ bm.GlobalNoUTurnSampler(
 )
 ```
 
+:::caution
+
+Functorch's [ahead of time (AOT) autograd compiler](https://pytorch.org/functorch/stable/aot_autograd.html) is used
+by default. If working with a non-static model or unexpected errors are encountered, you may need to manually
+disable the `nnc_compile` flag.
+
+:::
+
 The `GlobalNoUTurnSampler` has all the acceptance step size, covariance matrix, and acceptance probability tuning arguments of `GlobalHamiltonianMonteCarlo` as well as a few more parameters related to tuning the path length. While there are many optional parameters for this inference method, in practice, the parameters you are most likely to modify are `target_accept_prob` and `max_tree_depth`. When dealing with posteriors where the probability density has a more complicated shape, we benefit from taking smaller steps. Setting `target_accept_prob` to a higher value like `0.9` will lead to a more careful exploration of the space using smaller step sizes while still benefiting from some tuning of that step size. Since we will be taking smaller steps, we need to compensate by having a larger path length. This is accomplished by increasing `max_tree_depth`. Otherwise, using the defaults provided is highly recommended.
 
 A more complete explanation of parameters to `GlobalNoUTurnSampler` are provided below and in the [docs](https://beanmachine.org/api/beanmachine.ppl.html?highlight=nouturnsampler#beanmachine.ppl.GlobalNoUTurnSampler):
@@ -80,7 +87,7 @@ A more complete explanation of parameters to `GlobalNoUTurnSampler` are provided
 | `initial_step_size`    | The initial step size $\epsilon$ used in adaptive HMC. This value is simply the step size if tuning is disabled.                                                                                                                                          |
 | `multinomial_sampling` | Lets us decide between a faster multinomial sampler for the trajectory or the slice sampler described in the [original paper](https://arxiv.org/pdf/1111.4246.pdf). The option is useful for fairly comparing against other NUTS implementations.         |
 | `target_accept_prob`   | Indicates the acceptance probability which should be targeted by the step size tuning algorithm. While the optimal value is 65.1%, higher values have been show to be more robust leading to a default of 0.8.                                            |
-| `nnc_compile`          | NNC (Neural network compiler) is an experimental Pytorch JIT compiler that that transforms Pytorch programs to LLVM-compiled binaries. The model support is currently limited, so if your model fails, consider filing an issue and turning this flag off.|
+| `nnc_compile`          | NNC (neural network compiler) is a Pytorch JIT compiler that that transforms Pytorch programs to LLVM-compiled binaries. The model support is currently limited, so if your model fails, consider filing an issue and turning this flag off.|
 
 The parameters to `infer` are described below:
 

pyproject.toml

@@ -22,6 +22,8 @@ filterwarnings = [
     "default:distutils Version classes are deprecated.*:DeprecationWarning",
     # statsmodels imports a module that's deprecated since pandas 1.14.0
     "default:pandas.Int64Index is deprecated *:FutureWarning",
+    # functorch 0.1.0 imports deprecated _stateless module
+    "default:The `torch.nn.utils._stateless` code is deprecated*:DeprecationWarning",
 ]
 
 [tool.usort]

src/beanmachine/ppl/inference/compositional_infer.py

@@ -46,10 +46,11 @@ class _DefaultInference(BaseInference):
     Mixed inference class that handles both discrete and continuous RVs
     """
 
-    def __init__(self):
+    def __init__(self, nnc_compile: bool = True):
         self._disc_proposers = {}
         self._cont_proposer = None
         self._continuous_rvs = set()
+        self._nnc_compile = nnc_compile
 
     def get_proposers(
         self,
@@ -77,7 +78,10 @@ def get_proposers(
         else:
             if len(self._continuous_rvs):
                 continuous_proposer = NUTSProposer(
-                    world, self._continuous_rvs, num_adaptive_sample
+                    world,
+                    self._continuous_rvs,
+                    num_adaptive_sample,
+                    nnc_compile=self._nnc_compile,
                 )
                 self._cont_proposer = continuous_proposer
                 proposers.append(self._cont_proposer)
@@ -138,6 +142,7 @@ class CompositionalInference(BaseInference):
 
     Args:
         inference_dict: an optional inference configuration as shown above.
+        nnc_compile: where available, use NNC to compile proposers.
     """
 
     def __init__(
@@ -148,14 +153,15 @@ def __init__(
                 Union[BaseInference, Tuple[BaseInference, ...], EllipsisClass],
             ]
         ] = None,
+        nnc_compile: bool = True,
     ):
         self.config: Dict[Union[Callable, Tuple[Callable, ...]], BaseInference] = {}
         # create a set for the RV families that are being covered in the config; this is
         # useful in get_proposers to determine which RV needs to be handle by the
         # default inference method
         self._covered_rv_families = set()
 
-        default_inference = _DefaultInference()
+        default_inference = _DefaultInference(nnc_compile=nnc_compile)
         if inference_dict is not None:
             default_inference = inference_dict.pop(Ellipsis, default_inference)
             assert isinstance(default_inference, BaseInference)

src/beanmachine/ppl/inference/hmc_inference.py

@@ -86,6 +86,8 @@ class SingleSiteHamiltonianMonteCarlo(BaseInference):
         adapt_mass_matrix (bool): Whether to adapt the mass matrix. Defaults to True,
         target_accept_prob (float): Target accept prob. Increasing this value would lead
             to smaller step size. Defaults to 0.8.
+        nnc_compile: If True, NNC compiler will be used to accelerate the
+            inference.
     """
 
     def __init__(
@@ -95,7 +97,7 @@ def __init__(
         adapt_step_size: bool = True,
         adapt_mass_matrix: bool = True,
         target_accept_prob: float = 0.8,
-        nnc_compile: bool = False,
+        nnc_compile: bool = True,
     ):
         self.trajectory_length = trajectory_length
         self.initial_step_size = initial_step_size

src/beanmachine/ppl/inference/nuts_inference.py

@@ -34,8 +34,8 @@ class GlobalNoUTurnSampler(BaseInference):
             defaults to True.
         target_accept_prob (float): Target accept probability. Increasing this would
             lead to smaller step size. Defaults to 0.8.
-        nnc_compile: (Experimental) If True, NNC compiler will be used to accelerate the
-            inference (defaults to False).
+        nnc_compile: If True, NNC compiler will be used to accelerate the
+            inference.
     """
 
     def __init__(
@@ -47,7 +47,7 @@ def __init__(
         adapt_mass_matrix: bool = True,
         multinomial_sampling: bool = True,
         target_accept_prob: float = 0.8,
-        nnc_compile: bool = False,
+        nnc_compile: bool = True,
     ):
         self.max_tree_depth = max_tree_depth
         self.max_delta_energy = max_delta_energy

src/beanmachine/ppl/inference/proposer/hmc_proposer.py

@@ -45,8 +45,8 @@ class HMCProposer(BaseProposer):
         adapt_step_size: Flag whether to adapt step size, defaults to True.
         adapt_mass_matrix: Flat whether to adapt mass matrix, defaults to True.
         target_accept_prob: Target accept prob, defaults to 0.8.
-        nnc_compile: (Experimental) If True, NNC compiler will be used to accelerate the
-            inference (defaults to False).
+        nnc_compile: If True, NNC compiler will be used to accelerate the
+            inference.
     """
 
     def __init__(
@@ -59,7 +59,7 @@ def __init__(
         adapt_step_size: bool = True,
         adapt_mass_matrix: bool = True,
         target_accept_prob: float = 0.8,
-        nnc_compile: bool = False,
+        nnc_compile: bool = True,
     ):
         self.world = initial_world
         self._target_rvs = target_rvs

src/beanmachine/ppl/inference/proposer/nnc/__init__.py

@@ -26,7 +26,7 @@ def nnc_jit(
         # The setup code in `nnc.utils` will only be executed once in a Python session
         from beanmachine.ppl.inference.proposer.nnc.utils import nnc_jit as raw_nnc_jit
     except Exception as e:
-        logger.warn(
+        logger.warning(
             f"Fails to initialize NNC due to the following error: {str(e)}\n"
             "Falling back to default inference engine."
         )

src/beanmachine/ppl/inference/proposer/nuts_proposer.py

@@ -66,8 +66,8 @@ class NUTSProposer(HMCProposer):
         adapt_mass_matrix: Whether to adapt mass matrix using Welford Scheme, defaults to True.
         multinomial_sampling: Whether to use multinomial sampling as in [2], defaults to True.
         target_accept_prob: Target accept probability. Increasing this would lead to smaller step size. Defaults to 0.8.
-        nnc_compile: (Experimental) If True, NNC compiler will be used to accelerate the
-            inference (defaults to False).
+        nnc_compile: If True, NNC compiler will be used to accelerate the
+            inference.
     """
 
     def __init__(
@@ -82,7 +82,7 @@ def __init__(
         adapt_mass_matrix: bool = True,
         multinomial_sampling: bool = True,
         target_accept_prob: float = 0.8,
-        nnc_compile: bool = False,
+        nnc_compile: bool = True,
     ):
         # note that trajectory_length is not used in NUTS
         super().__init__(

src/beanmachine/ppl/inference/proposer/tests/hmc_proposer_test.py

@@ -30,7 +30,10 @@ def world():
 
 @pytest.fixture
 def hmc(world):
-    hmc_proposer = HMCProposer(world, world.latent_nodes, 10, trajectory_length=1.0)
+    # TODO(T130186904): fix NNC
+    hmc_proposer = HMCProposer(
+        world, world.latent_nodes, 10, trajectory_length=1.0, nnc_compile=False
+    )
     return hmc_proposer
 
 

src/beanmachine/ppl/inference/proposer/tests/nuts_proposer_test.py

@@ -30,7 +30,8 @@ def bar():
 def nuts():
     world = World(observations={bar(): torch.tensor(0.8)})
     world.call(bar())
-    nuts_proposer = NUTSProposer(world, world.latent_nodes, 10)
+    # TODO(T130186904): fix NNC
+    nuts_proposer = NUTSProposer(world, world.latent_nodes, 10, nnc_compile=False)
     return nuts_proposer
 
 
@@ -45,7 +46,10 @@ def tree_node(nuts):
 @pytest.fixture
 def tree_args(tree_node, nuts):
     initial_energy = nuts._hamiltonian(
-        nuts._positions, tree_node.momentums, nuts._mass_inv, nuts._pe
+        nuts._positions,
+        tree_node.momentums,
+        nuts._mass_inv,
+        nuts._pe,
     )
     return _TreeArgs(
         log_slice=-initial_energy,

src/beanmachine/ppl/inference/tests/compositional_infer_test.py

@@ -203,7 +203,7 @@ def test_block_inference_changing_support():
         {
             (model.K, model.component): bm.SingleSiteAncestralMetropolisHastings(),
             ...: bm.SingleSiteNewtonianMonteCarlo(),
-        }
+        },
     )
     sampler = compositional.sampler(queries, {}, num_samples=10, num_adaptive_samples=5)
     old_world = next(sampler)
@@ -221,8 +221,10 @@ def test_block_inference_changing_support():
             assert world[model.component(0)] is old_world[model.component(0)]
         old_world = world
 
+    # disable NNC because changing support => non-static model
     compositional = bm.CompositionalInference(
-        {(model.K, model.component): bm.SingleSiteAncestralMetropolisHastings()}
+        {(model.K, model.component): bm.SingleSiteAncestralMetropolisHastings()},
+        nnc_compile=False,
     )
     sampler = compositional.sampler(queries, {})
     with pytest.raises(KeyError):
@@ -240,7 +242,8 @@ def test_block_inference_changing_support():
 def test_block_inference_changing_shape():
     model = ChangingShapeModel()
     queries = [model.K()] + [model.component(j) for j in range(3)]
-    compositional = bm.CompositionalInference()
+    # disable NNC because changing shape => non-static model
+    compositional = bm.CompositionalInference(nnc_compile=False)
 
     # should run without error
     samples = compositional.infer(queries, {}, num_samples=5, num_chains=1).get_chain()