code: def test_reshape(self, preserve_dims, expected_output_shape):
      def f(inputs):
          return reshape.Reshape(
              output_shape=(-1, D), preserve_dims=preserve_dims
          )(inputs)
  
      init_fn, apply_fn = transform.transform(f)
      params = init_fn(None, jnp.ones([B, H, W, C, D]))
      outputs = apply_fn(params, None, np.ones([B, H, W, C, D]))
      self.assertEqual(outputs.shape, expected_output_shape)
  
imports: import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import reshape
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  from jax.experimental import jax2tf
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the ReshapeTest obj's underlying member method code with example usage: obj.test_reshape(preserve_dims, expected_output_shape) and returns: reshape