code: def where(condition, true_outputs, false_outputs):
      """Generalization of tf.where for nested structures.
  
      This generalization handles applying where across nested structures and the
      special case where the rank of the condition is smaller than the rank of the
      true and false cases.
  
      Args:
        condition: A boolean Tensor of shape [B, ...]. The shape of condition must
          be equal to or a prefix of the shape of true_outputs and false_outputs. If
          condition's rank is smaller than the rank of true_outputs and
          false_outputs, dimensions of size 1 are added to condition to make its
          rank match that of true_outputs and false_outputs in order to satisfy the
          requirements of tf.where.
        true_outputs: Tensor or nested tuple of Tensors of any dtype, each with
          shape [B, ...], to be split based on `condition`.
        false_outputs: Tensor or nested tuple of Tensors of any dtype, each with
          shape [B, ...], to be split based on `condition`.
  
      Returns:
        Interleaved output from `true_outputs` and `false_outputs` based on
        `condition`.
      """
      assert_same_structure(
          true_outputs,
          false_outputs,
          message='"true_outputs" and "false_outputs" structures do not match',
      )
  
      if tf.nest.flatten(true_outputs):
          condition_rank = tf.rank(condition)
  
          @tf.function  # allow-tf-function
          def per_field_where(t, f):
              tf.debugging.assert_rank_at_least(t, condition_rank)
              rank_difference = tf.rank(t) - condition_rank
              condition_shape = tf.concat(
                  [
                      tf.shape(condition),
                      tf.ones(rank_difference, dtype=tf.int32),
                  ],
                  axis=0,
              )
              per_field_condition = tf.reshape(
                  condition, condition_shape
              )
              return tf.compat.v2.where(per_field_condition, t, f)
  
          return tf.nest.map_structure(
              per_field_where, true_outputs, false_outputs
          )
      else:
          return tf.nest.map_structure(
              lambda t, f: tf.compat.v2.where(condition, t, f),
              true_outputs,
              false_outputs,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import numbers
  from typing import Optional, Text
  from absl import logging
  import numpy
  from six.moves import zip
  import tensorflow
  from tf_agents.typing import types
  from tf_agents.utils import composite
  import wrapt
  from tensorflow.python.util import nest
label_desc: Use image node_id 32 for calling a global function with example usage: where(condition, true_outputs, false_outputs) and returns: tf, tf, tf