import tensorflow as tf

      with client_graph.as_default():
        decoded_x = stage.decode(encoded_x, decode_params, shape=shape)
      with self.session(graph=client_graph):
        decoded_x = self.evaluate(decoded_x)

      return TestData(x, encoded_x, decoded_x, state, state_update_tensors,
                      updated_state)

    def _non_adaptive_one_to_many_encode_decode():
      """Implementation of the method for `EncodingStageInterface`."""
      server_graph = tf.Graph()
      with server_graph.as_default():
        x = input_fn()
        shape = py_utils.static_or_dynamic_shape(x)
        encode_params, decode_params = stage.get_params()
        encoded_x = stage.encode(x, encode_params)

      # Get all values out of TensorFlow as Python constants. This is a trivial
      # example of communication happening outside of TensorFlow.
      with self.session(graph=server_graph):
