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)
docstr: Constructs a :class:`Reshape` module.
  
  Args:
    output_shape: Shape to reshape the input tensor to while preserving its
      first ``preserve_dims`` dimensions. When the special value ``-1``
      appears in ``output_shape`` the corresponding size is automatically
      inferred. Note that ``-1`` can only appear once in ``output_shape``.
      To flatten all non-batch dimensions use :class:`Flatten`.
    preserve_dims: Number of leading dimensions that will not be reshaped.
      If negative, this is interpreted instead as the number of trailing
      dimensions to replace with the new shape.
    name: Name of the module.
  
  Raises:
    ValueError: If ``preserve_dims`` is zero.
docstr_tok: ["Constructs","a",":","class",":","`","Reshape","`","module",".","Args",":","output_shape",":","Shape","to","reshape","the","input","tensor","to","while","preserving","its","first","``","preserve_dims","``","dimensions",".","When","the","special","value","``","-1","``","appears","in","``","output_shape","``","the","corresponding","size","is","automatically","inferred",".","Note","that","``","-1","``","can","only","appear","once","in","``","output_shape","``",".","To","flatten","all","non-batch","dimensions","use",":","class",":","`","Flatten","`",".","preserve_dims",":","Number","of","leading","dimensions","that","will","not","be","reshaped",".","If","negative",",","this","is","interpreted","instead","as","the","number","of","trailing","dimensions","to","replace","with","the","new","shape",".","name",":","Name","of","the","module",".","Raises",":","ValueError",":","If","``","preserve_dims","``","is","zero","."]
code: def __init__(
      self,
      output_shape: Sequence[int],
      preserve_dims: int = 1,
      name: Optional[str] = None,
  ):
      """Constructs a :class:`Reshape` module.
  
      Args:
        output_shape: Shape to reshape the input tensor to while preserving its
          first ``preserve_dims`` dimensions. When the special value ``-1``
          appears in ``output_shape`` the corresponding size is automatically
          inferred. Note that ``-1`` can only appear once in ``output_shape``.
          To flatten all non-batch dimensions use :class:`Flatten`.
        preserve_dims: Number of leading dimensions that will not be reshaped.
          If negative, this is interpreted instead as the number of trailing
          dimensions to replace with the new shape.
        name: Name of the module.
  
      Raises:
        ValueError: If ``preserve_dims`` is zero.
      """
      super().__init__(name=name)
      if preserve_dims == 0:
          raise ValueError("Argument preserve_dims should be non-zero.")
      if output_shape.count(-1) > 1:
          raise ValueError("-1 can only occur once in `output_shape`.")
  
      self.output_shape = tuple(output_shape)
      self.preserve_dims = preserve_dims
  
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 Reshape object from inherited base classes: hk with example: obj = Reshape(output_shape, preserve_dims, name)