class_docstr: Flattens the input, preserving the batch dimension(s).
  
  By default, :class:`Flatten` combines all dimensions except the first.
  Additional leading dimensions can be preserved by setting ``preserve_dims``.
  
  >>> x = jnp.ones([3, 2, 4])
  >>> flat = hk.Flatten()
  >>> flat(x).shape
  (3, 8)
  
  When the input to flatten has fewer than ``preserve_dims`` dimensions it is
  returned unchanged:
  
  >>> x = jnp.ones([3])
  >>> flat(x).shape
  (3,)
  
  Alternatively, a negative value of `preserve_dims` specifies the number of
  trailing dimensions flattened:
  
  >>> x = jnp.ones([3, 2, 4])
  >>> negative_flat = hk.Flatten(preserve_dims=-2)
  >>> negative_flat(x).shape
  (3, 8)
  
  This allows the same module to be seamlessly applied to a single element or a
  batch of elements with the same element shape:
  
  >> negative_flat(x[0]).shape
  (8,)
code: def __init__(
      self,
      preserve_dims: int = 1,
      name: Optional[str] = None,
  ):
      super().__init__(
          output_shape=(-1,), preserve_dims=preserve_dims, name=name
      )
  
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 1 to create a new Flatten object from inherited base classes: Reshape with example: obj = Flatten(preserve_dims, name)