code: def testFlatten(self, batch_size):
      in_shape = [2, 3, 4, 5]
      inputs = tf.ones([batch_size] + in_shape)
      mod = reshape.Flatten()
      output = mod(inputs)
      flattened_size = np.prod(in_shape, dtype=int)
      self.assertEqual(output.shape, [batch_size, flattened_size])
  
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 1 for calling the FlattenTest obj's underlying member method code with example usage: obj.testFlatten(batch_size) without return types