docs: def _build(self, inputs):
      """Connects the module into the graph, with input Tensor `inputs`.
  
      Args:
        inputs: A Tensor of shape [b_1, b_2, ..., b_preserve_dims,
                                   b_preserve_dims+1, ...].
  
      Returns:
        A Tensor of shape [b_1, b_2, ..., b_preserve_dims,
                           b_reshape_1, b_reshape_2, ...],
          with reshaping defined by the constructor `shape` parameter.
  
      Raises:
        ValueError: If output shape is incompatible with input shape; or if
            shape array contains non numeric entries; or if shape array contains
            more than 1 wildcard -1; or if the input array contains unknown,
            non-preserved dimensions (except when the unknown dimension is the
            only non-preserved dimension and doesn't actually need reshaping).
      """
      full_input_shape = inputs.get_shape().as_list()
      if len(full_input_shape) < self._preserve_dims:
        raise ValueError("Input tensor has {} dimensions, should have at least "
                         "as many as preserve_dims={}".format(
                             len(full_input_shape),
                             self._preserve_dims))
      self._input_shape = full_input_shape[self._preserve_dims:]
  
      if callable(self._shape):
        self._shape = tuple(self._shape())
  
      # Special-case of 1 non-preserved dimension, where no reshape is necessary.
      # This is useful if the non-preserved dimension of `inputs` is unknown
      # at build time.
      if len(self._input_shape) == 1 and len(self._shape) == 1:
        if self._shape[0] == -1 or self._shape[0] == self._input_shape[0]:
          return inputs
        else:
          if self._input_shape[0] is None:
            raise ValueError("Unknown non-preserved dimensions are not allowed "
                             "in the input to BatchReshape unless it is only one "
                             "and the desired shape is (-1,).")
          else:
            raise ValueError("Output shape is incompatible with input shape")
  
      if not all([isinstance(x, numbers.Integral) and (x > 0 or x == -1)
                  for x in self._shape]):
        raise ValueError(
            "Desired shape can only contain positive integral numbers "
            "and the wildcard -1. Given shape {}".format(self._shape))
  
      if self._shape.count(-1) > 1:
        raise ValueError(
            "Wildcard -1 can appear only once in desired output shape. "
            "Given shape {}".format(self._shape))
  
      preserved_shape = tf.shape(inputs)[:self._preserve_dims]
      # Slicing the shape tensor loses information, we keep it in a list.
      preserved_shape_list = inputs.get_shape()[:self._preserve_dims]
  
      # Except in the case above where no reshape is needed, we do not allow
      # unknown non-preserved dimensions in the input.
      if None in self._input_shape:
        raise ValueError("Unknown non-preserved dimensions are not allowed in "
                         "the input to BatchReshape unless it is only one and the"
                         " desired shape is (-1,). The offending non-preserved "
                         "input shape is {}".format(self._input_shape))
      if self._shape.count(-1) > 0:
        trailing_shape = self._infer_shape(self._input_shape)
      else:
        trailing_shape = self._shape
  
      if np.prod(self._input_shape) != np.prod(trailing_shape):
        raise ValueError("Output shape is incompatible with input shape")
  
      shape = tf.concat([preserved_shape, trailing_shape], 0)
      output = tf.reshape(inputs, shape)
  
      # Include shape information that was lost when we sliced the shape tensor.
      shape_list = preserved_shape_list.concatenate(trailing_shape)
      output.set_shape(shape_list)
  
      return output