code: def testReverse(self, batch_size):
      mod = reshape.Flatten(preserve_dims=4)
      inputs = tf.ones([batch_size, 5, 84, 84, 3, 2])
      output = mod(inputs)
      self.assertEqual(output.shape, inputs.shape.as_list()[:4] + [6])
      mod_r = mod.reversed()
      output_r = mod_r(output)
      self.assertEqual(output_r.shape, inputs.shape)
  
imports: from absl.testing import parameterized
  import numpy
  from sonnet.src import reshape
  from sonnet.src import test_utils
  import tensorflow
label_desc: Use image node_id 9 for calling the FlattenTest obj's underlying member method code with example usage: obj.testReverse(batch_size) without return types