code: def _generalized_where(mask, value, old_value):
      """Version of tf1.where that broadcasts `value` to `old_value`."""
      mask = tf.convert_to_tensor(value=mask, name="mask")
      mask.shape.assert_has_rank(1)
      value = tf.convert_to_tensor(
          value=value, name="value", dtype=old_value.dtype
      )
      if (
          not value.shape.is_fully_defined()
          or not old_value.shape.is_fully_defined()
          or value.shape != old_value.shape
      ):
          # We force broadcast value w/ current, e.g. for program constants.
          if old_value.dtype == tf.bool:
              value |= tf.zeros_like(old_value)
          else:
              value += tf.zeros_like(old_value)
      new_value = tf1.where(mask, value, old_value, name="new_value")
      # TODO(b/78655271): Do we need 'new_val.set_shape(old_value.shape)'?
      return new_value
  
imports: import collections
  import contextlib
  import six
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import xla
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 2 for calling a global function with example usage: _generalized_where(mask, value, old_value) and returns: new_value