Implementing Uncorrected Hamiltonian Annealing / Differentiable Annealed Importance Sampling

martinjankowiak · December 6, 2022, 1:04am

afaik the place to start would be initialize_model

pyro-ppl/pyro/blob/6d61ea94ee7c4eef7aab369e9f2f39e7af82c98e/pyro/infer/mcmc/util.py#L370


      
                          if num_chains == 1:
                              return {k: v[0] for k, v in params_per_chain.items()}
                          else:
                              return {k: torch.stack(v) for k, v in params_per_chain.items()}
                  trace = None
              raise ValueError(
                  "Model specification seems incorrect - cannot find valid initial params."
              )
          
          

          
def initialize_model(
              model,
              model_args=(),
              model_kwargs={},
              transforms=None,
              max_plate_nesting=None,
              jit_compile=False,
              jit_options=None,
              skip_jit_warnings=False,
              num_chains=1,
              init_strategy=init_to_uniform,

although @fehiepsi might have a better suggestion. note it might be a bit complicated to get all the details right, given the complexity of transforms and the like.

also note that any pyro implementation of uha/dais is expected to be somewhat slow compared to numpyro. basically because there will be overhead from a python for loop, which jax can compile away. also pytorch isn’t great about optimizing gradients. this is one of the reasons we implemented these methods in numpyro.