docs: def _infer_shape(self, dimensions):
      """Replaces the -1 wildcard in the output shape vector.
  
      This function infers the correct output shape given the input dimensions.
  
      Args:
        dimensions: List of input non-batch dimensions.
  
      Returns:
        Tuple of non-batch output dimensions.
      """
      # Size of input
      n = np.prod(dimensions)
      # Size of output where defined
      m = np.prod(abs(np.array(self._shape)))
      # Replace wildcard
      v = np.array(self._shape)
      v[v == -1] = n // m
      return tuple(v)