How to perform ADVI only on some of the variables in a model?

Questions v5
7

Let’s say now I relaxed the condition. Now model M1 should be trained from the end point of model M2. Would the below procedure work?

with M2:
        
  inference = pm.ADVI(random_seed=seed)

  pm.set_data({"ann_input": u_scale_train[index_train]})
  tracker = pm.callbacks.Tracker(
  mean= inference.approx.mean.eval,  # callable that returns mean
  std= inference.approx.std.eval  # callable that returns std
  )

  approx = pm.fit(n= 10000, random_seed=seed, callbacks=[pm.callbacks.CheckParametersConvergence(tolerance=1e-2)], method=inference, obj_optimizer=adam(learning_rate=0.05,decay_iter=4000))
  idata = approx.sample(2000,random_seed=seed) 

with M1:
  #inference = pm.ADVI(random_seed=seed)

  pm.set_data({"ann_input": u_scale_train[index_train]})
  tracker = pm.callbacks.Tracker(
  mean= inference.approx.mean.eval,  # callable that returns mean
  std= inference.approx.std.eval  # callable that returns std
  )

  pm.ADVI.refine(inference,n=4000)
  idata_servicer = approx.sample(2000,random_seed=seed)

In the above procedure, the parameters (p2) of the “M2” get trained and in the second part, I use the same set of parameters and refine the “M1” model which contains some additional parameters (p2+p1). In this procedure, some parameters (p2) of the “M1” are getting trained but not from the very start and I other parameters (p1) are trained from the start. Would the above code work to solve the relaxed problem?
I think I missing something because parameters “p1” are not getting updated.