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).

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