message: Fix flaky mixed_precision/experimental:keras_test.
diff: @@ -81,6 +81,11 @@ def create_identity_with_nan_gradients_fn(have_nan_gradients):
  """Function whose gradient is NaN iff `have_nan_gradients` is True."""
  x = array_ops.identity(x)
  def grad(dx):
  + # We need this control dependency, because otherwise the NaN could be
  + # produced before `dx`. This in turn could cause the final gradient to be
  + # produced because `dx`, causing the loss scale to be updated before `dx`,
  + # which can cause `tf.assert_equal`s to fail.
  + with ops.control_dependencies([dx]):
  nan_scalar = constant_op.constant(float('NaN'), dtype=dx.dtype)
  return control_flow_ops.cond(
  have_nan_gradients,
  