class_docstr: Reshapes input Tensor, preserving the batch dimension.
  
  For example, given an input tensor with shape ``[B, H, W, C, D]``::
  
      >>> B, H, W, C, D = range(1, 6)
      >>> x = jnp.ones([B, H, W, C, D])
  
  The default behavior when ``output_shape`` is ``(-1, D)`` is to flatten
  all dimensions between ``B`` and ``D``::
  
      >>> mod = hk.Reshape(output_shape=(-1, D))
      >>> assert mod(x).shape == (B, H*W*C, D)
  
  You can change the number of preserved leading dimensions via
  ``preserve_dims``::
  
      >>> mod = hk.Reshape(output_shape=(-1, D), preserve_dims=2)
      >>> assert mod(x).shape == (B, H, W*C, D)
  
      >>> mod = hk.Reshape(output_shape=(-1, D), preserve_dims=3)
      >>> assert mod(x).shape == (B, H, W, C, D)
  
      >>> mod = hk.Reshape(output_shape=(-1, D), preserve_dims=4)
      >>> assert mod(x).shape == (B, H, W, C, 1, D)
  
  Alternatively, a negative value of ``preserve_dims`` specifies
  the number of trailing dimensions to replace with ``output_shape``::
  
      >>> mod = hk.Reshape(output_shape=(-1, D), preserve_dims=-3)
      >>> assert mod(x).shape == (B, H, W*C, D)
  
  This is useful in the case of applying the same module to batched
  and unbatched outputs::
  
      >>> mod = hk.Reshape(output_shape=(-1, D), preserve_dims=-3)
      >>> assert mod(x[0]).shape == (H, W*C, D)
code: def __call__(self, inputs):
      if inputs.ndim <= self.preserve_dims:
          return inputs
  
      if -1 in self.output_shape:
          reshaped_shape = _infer_shape(
              self.output_shape, inputs.shape[self.preserve_dims :]
          )
      else:
          reshaped_shape = self.output_shape
      shape = inputs.shape[: self.preserve_dims] + reshaped_shape
      return jnp.reshape(inputs, shape)
  
file_path: dm-haiku/haiku/_src/reshape.py
imports: from collections.abc import Sequence
  from typing import Optional
  from haiku._src import module
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the Reshape obj's underlying member method code with example usage: obj.__call__(inputs) and returns: jnp, inputs