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