Fix bug in particle filter where static evaluation could block the gradient signal.

PiperOrigin-RevId: 384310858

Author: davmre

tensorflow_probability/python/experimental/mcmc/particle_filter.py

@@ -500,7 +500,7 @@ def _compute_observation_log_weights(step,
         lambda x, step=step: tf.gather(x, observation_idx), observations)
 
     log_weights = observation_fn(step, particles).log_prob(observation)
-    return ps.where(step_has_observation,
+    return tf.where(step_has_observation,
                     log_weights,
                     tf.zeros_like(log_weights))
 

tensorflow_probability/python/experimental/mcmc/particle_filter_test.py

@@ -561,6 +561,23 @@ def transition_fn_no_batch_shape(_, particles):
           run_filter(transition_fn=valid_transition_fn,
                      proposal_fn=transition_fn_no_batch_shape))
 
+  @test_util.jax_disable_test_missing_functionality('Gradient of while_loop.')
+  def test_marginal_likelihood_gradients_are_defined(self):
+
+    def marginal_log_likelihood(level_scale, noise_scale):
+      _, _, _, lps = tfp.experimental.mcmc.particle_filter(
+          observations=tf.convert_to_tensor([1., 2., 3., 4., 5.]),
+          initial_state_prior=tfd.Normal(loc=0, scale=1.),
+          transition_fn=lambda _, x: tfd.Normal(loc=x, scale=level_scale),
+          observation_fn=lambda _, x: tfd.Normal(loc=x, scale=noise_scale),
+          num_particles=4,
+          seed=test_util.test_seed())
+      return tf.reduce_sum(lps)
+
+    _, grads = tfp.math.value_and_gradient(marginal_log_likelihood, 1.0, 1.0)
+    self.assertAllNotNone(grads)
+    self.assertNotAllZero(grads)
+
 
 # TODO(b/186068104): add tests with dynamic shapes.
 class ParticleFilterTestFloat32(_ParticleFilterTest):

tensorflow_probability/python/experimental/mcmc/sequential_monte_carlo_kernel.py

@@ -277,7 +277,7 @@ def one_step(self, state, kernel_results, seed=None):
         (resampled_particles,
          resample_indices,
          log_weights) = tf.nest.map_structure(
-             lambda r, p: ps.where(do_resample, r, p),
+             lambda r, p: tf.where(do_resample, r, p),
              (resampled_particles, resample_indices, uniform_weights),
              (state.particles, _dummy_indices_like(resample_indices),
               normalized_log_weights))