FFJORD Trace Computation during Training vs. Evaluation

After training the model in the FFJORD demo: https://www.tensorflow.org/probability/examples/FFJORD_Demo, I use a checkpoint manager to save and load:

solver = tfp.math.ode.DormandPrince(atol=1e-5)
ode_solve_fn = solver.solve
trace_augmentation_fn = tfb.ffjord.trace_jacobian_hutchinson

bijectors = []

for _ in range(STACKED_FFJORDS):
  mlp_model = MLP_ODE(NUM_HIDDEN, NUM_LAYERS, NUM_OUTPUT)
  next_ffjord = tfb.FFJORD(
      state_time_derivative_fn=mlp_model,ode_solve_fn=ode_solve_fn,
      trace_augmentation_fn=trace_augmentation_fn)
  bijectors.append(next_ffjord)

stacked_ffjord = tfb.Chain(bijectors[::-1])
transformed_distribution = tfd.TransformedDistribution(
    distribution=base_distribution, bijector=stacked_ffjord)



learning_rate = tf.Variable(LR, trainable=False)
optimizer = tf.keras.optimizers.Adam(learning_rate)
checkpoint = tf.train.Checkpoint(step=tf.Variable(1),optimizer=optimizer, model=transformed_distribution)
manager = tf.train.CheckpointManager(
                checkpoint, directory="./models/", max_to_keep=5)

I would like to change the trace computation after training to use tfb.ffjord.trace_jacobian_exact instead of Hutchinson, but I can't load the model if I change the trace. Any ideas on how to do this?

Thanks

Start with the FFJORD demo and the checkpoint construction shown in the issue. Compare loading behavior when using trace_jacobian_hutchinson versus trace_jacobian_exact, then verify that a checkpoint created during training can be restored after the trace function is changed for evaluation.