docstr: Tests that encoder works with stages that need input shape for decode.
  
  This test chains two stages with this property.
docstr_tok: ["Tests","that","encoder","works","with","stages","that","need","input","shape","for","decode",".","This","test","chains","two","stages","with","this","property","."]
code: def test_decode_needs_input_shape(self):
      """Tests that encoder works with stages that need input shape for decode.
  
      This test chains two stages with this property.
      """
      encoder = (
          core_encoder.EncoderComposer(
              test_utils.ReduceMeanEncodingStage()
          )
          .add_parent(test_utils.ReduceMeanEncodingStage(), RM_VALS)
          .make()
      )
      x = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])
      encode_params, decode_params = encoder.get_params(
          encoder.initial_state()
      )
      encoded_x, _, input_shapes = encoder.encode(x, encode_params)
      decoded_x = encoder.decode(encoded_x, decode_params, input_shapes)
      encoded_x, decoded_x = self.evaluate([encoded_x, decoded_x])
  
      self.assertAllEqual([3.0] * 5, decoded_x)
      self.assertAllEqual({RM_VALS: {RM_VALS: 3.0}}, encoded_x)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import numpy
  import tensorflow
  from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import core_encoder
  from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import encoding_stage
  from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing import test_utils
label_desc: Use image node_id 9 for calling the EncoderTest obj's underlying member method code with example usage: obj.test_decode_needs_input_shape() without return types