jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, name, inputs, tower_setup, n_convs=2, n_features=None, dilations=None, strides=None,
               filter_size=None, activation="relu", dropout=0.0, batch_norm_decay=BATCH_NORM_DECAY_DEFAULT,
               l2=L2_DEFAULT):
    super(ResidualUnit2, self).__init__()
    curr, n_features_inp = prepare_input(inputs)
    res = curr
    assert n_convs >= 1, n_convs

    if dilations is not None:
      assert strides is None
    elif strides is None:
      strides = [[1, 1]] * n_convs
    if filter_size is None:
      filter_size = [[3, 3]] * n_convs
    if n_features is None:
      n_features = n_features_inp
    if not isinstance(n_features, list):
      n_features = [n_features] * n_convs

    with tf.variable_scope(name):
      curr = self.create_and_apply_batch_norm(curr, n_features_inp, batch_norm_decay, tower_setup, "bn0")
      curr = get_activation(activation)(curr)
      curr = apply_dropout(curr, dropout)

      if strides is None:
        strides_res = [1, 1]
      else:
        strides_res = numpy.prod(strides, axis=0).tolist()
      if (n_features[-1] != n_features_inp) or (strides_res != [1, 1]):
        W0 = self.create_weight_variable("W0", [1, 1] + [n_features_inp, n_features[-1]], l2, tower_setup)
        if dilations is None:
          res = conv2d(curr, W0, strides_res)
        else:
          res = conv2d(curr, W0)

      W1 = self.create_weight_variable("W1", filter_size[0] + [n_features_inp, n_features[0]], l2, tower_setup)
      if dilations is None:
        curr = conv2d(curr, W1, strides[0])
      else:
        curr = conv2d_dilated(curr, W1, dilations[0])
      for idx in range(1, n_convs):
        curr = self.create_and_apply_batch_norm(curr, n_features[idx - 1], batch_norm_decay,
                                                tower_setup, "bn" + str(idx + 1))
        curr = get_activation(activation)(curr)
        Wi = self.create_weight_variable("W" + str(idx + 1), filter_size[idx] + [n_features[idx - 1], n_features[idx]],
                                         l2, tower_setup)
        if dilations is None:
          curr = conv2d(curr, Wi, strides[idx])
        else:
          curr = conv2d_dilated(curr, Wi, dilations[idx])

    curr += res
    self.outputs = [curr]
//  [python]
def _generate_cells(dim):
    ele_dim = [d - 1 for d in dim if d > 1]
    ele_no = n.prod(ele_dim, dtype=int)
    spatial_dim = len(ele_dim)

    if spatial_dim == 1:
        # cells are lines in 1D
        cells = numpy.empty((ele_no, 3), dtype=int)
        cells[:, 0] = 2
        cells[:, 1] = numpy.arange(ele_no, dtype=int)
        cells[:, 2] = cells[:, 1] + 1
        cell_types = numpy.full(ele_no, 3, dtype=int)

    elif spatial_dim == 2:
        # cells are quad in 2D
        cells = numpy.empty((ele_no, 5), dtype=int)
        cells[:, 0] = 4
        cells[:, 1] = numpy.arange(0, ele_no, dtype=int)
        cells[:, 1] += numpy.arange(0, ele_no, dtype=int) // ele_dim[0]
        cells[:, 2] = cells[:, 1] + 1
        cells[:, 3] = cells[:, 1] + 2 + ele_dim[0]
        cells[:, 4] = cells[:, 3] - 1
        cell_types = numpy.full(ele_no, 9, dtype=int)
    else:
        # cells are hex in 3D
        cells = numpy.empty((ele_no, 9), dtype=int)
        cells[:, 0] = 8
        cells[:, 1] = numpy.arange(ele_no)
        cells[:, 1] += (ele_dim[0] + ele_dim[1] + 1) * (
            numpy.arange(ele_no) // (ele_dim[0] * ele_dim[1])
        )
        cells[:, 1] += (numpy.arange(ele_no) % (ele_dim[0] * ele_dim[1])) // ele_dim[0]
        cells[:, 2] = cells[:, 1] + 1
        cells[:, 3] = cells[:, 1] + 2 + ele_dim[0]
        cells[:, 4] = cells[:, 3] - 1
        cells[:, 5] = cells[:, 1] + (1 + ele_dim[0]) * (1 + ele_dim[1])
        cells[:, 6] = cells[:, 5] + 1
        cells[:, 7] = cells[:, 5] + 2 + ele_dim[0]
        cells[:, 8] = cells[:, 7] - 1
        cell_types = numpy.full(ele_no, 12, dtype=int)

    return cells.reshape(-1), cell_types
//  [python]
def _generate_cells(dim):
    ele_dim = [d - 1 for d in dim if d > 1]
    ele_no = numpy.prod(ele_dim, dtype=int)
    spatial_dim = len(ele_dim)

    if spatial_dim == 1:
        # cells are lines in 1D
        cells = numpy.empty((ele_no, 3), dtype=int)
        cells[:, 0] = 2
        cells[:, 1] = numpy.arange(ele_no, dtype=int)
        cells[:, 2] = cells[:, 1] + 1
        cell_types = numpy.full(ele_no, 3, dtype=int)

    elif spatial_dim == 2:
        # cells are quad in 2D
        cells = numpy.empty((ele_no, 5), dtype=int)
        cells[:, 0] = 4
        cells[:, 1] = numpy.arange(0, ele_no, dtype=int)
        cells[:, 1] += numpy.arange(0, ele_no, dtype=int) // ele_dim[0]
        cells[:, 2] = cells[:, 1] + 1
        cells[:, 3] = cells[:, 1] + 2 + ele_dim[0]
        cells[:, 4] = cells[:, 3] - 1
        cell_types = numpy.full(ele_no, 9, dtype=int)
    else:
        # cells are hex in 3D
        cells = numpy.empty((ele_no, 9), dtype=int)
        cells[:, 0] = 8
        cells[:, 1] = numpy.arange(ele_no)
        cells[:, 1] += (ele_dim[0] + ele_dim[1] + 1) * (
            numpy.arange(ele_no) // (ele_dim[0] * ele_dim[1])
        )
        cells[:, 1] += (numpy.arange(ele_no) % (ele_dim[0] * ele_dim[1])) // ele_dim[0]
        cells[:, 2] = cells[:, 1] + 1
        cells[:, 3] = cells[:, 1] + 2 + ele_dim[0]
        cells[:, 4] = cells[:, 3] - 1
        cells[:, 5] = cells[:, 1] + (1 + ele_dim[0]) * (1 + ele_dim[1])
        cells[:, 6] = cells[:, 5] + 1
        cells[:, 7] = cells[:, 5] + 2 + ele_dim[0]
        cells[:, 8] = cells[:, 7] - 1
        cell_types = numpy.full(ele_no, 12, dtype=int)

    return cells.reshape(-1), cell_types
//  [python]
def _find_best_partition(dims: Union[List[int], np.ndarray]) -> int:
  """
  Find the most-levelled partition of `dims`.
  A levelled partitioning is a partitioning such that
  np.prod(dim[:partition]) and np.prod(dim[partition:])
  are as close as possible.
  Args:
    dims: A list or np.ndarray of integers.
  Returns:
    int: The best partitioning.
  """
  if len(dims) == 1:
    raise ValueError(
        'expecting dims with a length of at least 2, got len(dims) =1')
  diffs = [
      np.abs(reduce(lambda x, y: x * y, dims[0:n]) - reduce(lambda x, y: x * y, dims[n::]))
      for n in range(1, len(dims))
  ]
  min_inds = np.nonzero(diffs == np.min(diffs))[0]
  if len(min_inds) > 1:
    right_dims = [reduce(lambda x, y: x * y, dims[min_ind + 1:]) for min_ind in min_inds]
    min_ind = min_inds[np.argmax(right_dims)]
  else:
    min_ind = min_inds[0]
  return min_ind + 1
//  [python]
def _find_best_partition(dims: Union[List[int], np.ndarray]) -> int:
  """
  Find the most-levelled partition of `dims`.
  A levelled partitioning is a partitioning such that
  np.prod(dim[:partition]) and np.prod(dim[partition:])
  are as close as possible.
  Args:
    dims: A list or np.ndarray of integers.
  Returns:
    int: The best partitioning.
  """
  if len(dims) == 1:
    raise ValueError(
        'expecting dims with a length of at least 2, got len(dims) =1')
  diffs = [
      np.abs(np.prod(dims[0:n]) - np.prod(dims[n::]))
      for n in range(1, len(dims))
  ]
  min_inds = np.nonzero(diffs == np.min(diffs))[0]
  if len(min_inds) > 1:
    right_dims = [np.prod(dims[min_ind + 1:]) for min_ind in min_inds]
    min_ind = min_inds[np.argmax(right_dims)]
  else:
    min_ind = min_inds[0]
  return min_ind + 1
//  [python]
def get_required_memory_size(self, view_model):
        # type: (BrainViewerModel) -> numpy.ndarray
        """
        Assume one page doesn't get 'dumped' in time and it is highly probably that
        two consecutive pages will be in the same time in memory.
        """
        time_series = self.load_entity_by_gid(view_model.time_series)
        overall_shape = time_series.get_data_shape()
        used_shape = (overall_shape[0] / (self.PAGE_SIZE * 2.0), overall_shape[1], overall_shape[2], overall_shape[3])
        return numpy.multiply.reduce(used_shape) * 8.0
//  [python]
def get_required_memory_size(self, view_model):
        # type: (BrainViewerModel) -> numpy.ndarray
        """
        Assume one page doesn't get 'dumped' in time and it is highly probably that
        two consecutive pages will be in the same time in memory.
        """
        time_series = self.load_entity_by_gid(view_model.time_series)
        overall_shape = time_series.get_data_shape()
        used_shape = (overall_shape[0] / (self.PAGE_SIZE * 2.0), overall_shape[1], overall_shape[2], overall_shape[3])
        return numpy.prod(used_shape) * 8.0
//  [python]
def _cartesian(arraySizes, out=None):
    """
    Generate a cartesian product of input arrays.

    Parameters
    ----------
    arraySizes : list
        list of size of arrays
    out : numpy.ndarray, optional
        Array to place the cartesian product in.

    Returns
    -------
    numpy.ndarray
        2-D array of shape (product(arraySizes), len(arraySizes)) containing cartesian products

    Notes
    -----
    -  2016-06-02 - Obtained from http://stackoverflow.com/questions/1208118/using-numpy-to-build-an-array-of-all-combinations-of-two-arrays
    """
    arrays = []
    for i in range(len(arraySizes)):
        arrays.append(numpy.arange(0, arraySizes[i]))

    arrays = [numpy.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

    n = reduce(lambda x, y: x * y, [x.size for x in arrays])
    if out is None:
        out = numpy.zeros([n, len(arrays)], dtype=dtype)

    m = n // arrays[0].size
    out[:, 0] = numpy.repeat(arrays[0], m)
    if arrays[1:]:
        _cartesian(arraySizes[1:], out=out[0:m, 1:])
        for j in range(1, arrays[0].size):
            out[j * m : (j + 1) * m, 1:] = out[0:m, 1:]
    return out
//  [python]
def _cartesian(arraySizes, out=None):
    """
    Generate a cartesian product of input arrays.

    Parameters
    ----------
    arraySizes : list
        list of size of arrays
    out : numpy.ndarray, optional
        Array to place the cartesian product in.

    Returns
    -------
    numpy.ndarray
        2-D array of shape (product(arraySizes), len(arraySizes)) containing cartesian products

    Notes
    -----
    -  2016-06-02 - Obtained from http://stackoverflow.com/questions/1208118/using-numpy-to-build-an-array-of-all-combinations-of-two-arrays
    """
    arrays = []
    for i in range(len(arraySizes)):
        arrays.append(numpy.arange(0, arraySizes[i]))

    arrays = [numpy.asarray(x) for x in arrays]
    dtype = arrays[0].dtype

    n = numpy.prod([x.size for x in arrays])
    if out is None:
        out = numpy.zeros([n, len(arrays)], dtype=dtype)

    m = n // arrays[0].size
    out[:, 0] = numpy.repeat(arrays[0], m)
    if arrays[1:]:
        _cartesian(arraySizes[1:], out=out[0:m, 1:])
        for j in range(1, arrays[0].size):
            out[j * m : (j + 1) * m, 1:] = out[0:m, 1:]
    return out
//  [python]
def _construct_collection(
    orders,
    dist,
    x_lookup,
    w_lookup,
):
    """Create a collection of {abscissa: weight} key-value pairs."""
    order = numpy.min(orders)
    skew = orders - order

    # Indices and coefficients used in the calculations
    indices = numpoly.glexindex(order - len(dist) + 1, order + 1, dimensions=len(dist))
    coeffs = numpy.sum(indices, -1)
    coeffs = (2 * ((order - coeffs + 1) % 2) - 1) * comb(len(dist) - 1, order - coeffs)

    collection = defaultdict(float)
    for bidx, coeff in zip(indices + skew, coeffs.tolist()):
        abscissas = [value[idx] for idx, value in zip(bidx, x_lookup)]
        weights = [value[idx] for idx, value in zip(bidx, w_lookup)]
        for abscissa, weight in zip(product(*abscissas), product(*weights)):
            collection[abscissa] += numpy.prod(weight) * coeff

    return collection
//  [python]
def count_network_parameters(model):

    parameters = filter(lambda p: p.requires_grad, model.parameters())
    N = sum([n.prod(p.size()) for p in parameters])

    return N
//  [python]
def count_network_parameters(model):

    parameters = filter(lambda p: p.requires_grad, model.parameters())
    N = sum([np.prod(p.size()) for p in parameters])

    return N
//  [python]
def get_binary_tensors(binary_paths: List[str], info: AppInputInfo, batch_sizes: List[int]) -> List[Tensor]:
    num_shapes = len(info.shapes)
    num_binaries = len(binary_paths)
    niter = max(num_shapes, num_binaries)
    processed_frames = 0
    tensors = []
    for i in range(niter):
        shape_id = i % num_shapes
        dtype = get_dtype(info.element_type)
        shape = list(info.shapes[shape_id])
        binaries = np.ndarray(shape=shape, dtype=dtype)
        binary_index = processed_frames
        current_batch_size = batch_sizes[shape_id]
        for b in range(current_batch_size):
            binary_index %= num_binaries
            binary_filename: str = binary_paths[binary_index]
            extension = binary_filename.lower().split('.')[-1]
            if extension == "bin":
                binary_file_size = os.path.getsize(binary_filename)
                blob_size = dtype.itemsize * int(np.multiply.reduce(shape))
                if blob_size != binary_file_size:
                    raise Exception(
                        f"File {binary_filename} contains {binary_file_size} bytes but model expects {blob_size}")
                from_file = np.reshape(np.fromfile(binary_filename, dtype), shape)
                if info.layout.has_name("N"):
                    binaries[[None] * info.layout.get_index_by_name("N") + [b]] = from_file
                else:
                    binaries = from_file
            else:
                raise Exception(
                    f"Unsupported binary file type: {extension}")

            binary_index += 1
        processed_frames += current_batch_size
        tensors.append(Tensor(binaries))
    return tensors
//  [python]
def get_binary_tensors(binary_paths: List[str], info: AppInputInfo, batch_sizes: List[int]) -> List[Tensor]:
    num_shapes = len(info.shapes)
    num_binaries = len(binary_paths)
    niter = max(num_shapes, num_binaries)
    processed_frames = 0
    tensors = []
    for i in range(niter):
        shape_id = i % num_shapes
        dtype = get_dtype(info.element_type)
        shape = list(info.shapes[shape_id])
        binaries = np.ndarray(shape=shape, dtype=dtype)
        binary_index = processed_frames
        current_batch_size = batch_sizes[shape_id]
        for b in range(current_batch_size):
            binary_index %= num_binaries
            binary_filename: str = binary_paths[binary_index]
            extension = binary_filename.lower().split('.')[-1]
            if extension == "bin":
                binary_file_size = os.path.getsize(binary_filename)
                blob_size = dtype.itemsize * int(np.prod(shape))
                if blob_size != binary_file_size:
                    raise Exception(
                        f"File {binary_filename} contains {binary_file_size} bytes but model expects {blob_size}")
                from_file = np.reshape(np.fromfile(binary_filename, dtype), shape)
                if info.layout.has_name("N"):
                    binaries[[None] * info.layout.get_index_by_name("N") + [b]] = from_file
                else:
                    binaries = from_file
            else:
                raise Exception(
                    f"Unsupported binary file type: {extension}")

            binary_index += 1
        processed_frames += current_batch_size
        tensors.append(Tensor(binaries))
    return tensors
//  [python]
def __init__(self, which_set, multi_target=False, stop=None):
        """
        parameters
        ----------

        which_set : str
            Must be 'train' or 'test'.

        multi_target : bool
            If False, each label is an integer labeling the image catergory. If
            True, each label is a vector: [category, instance, lighting,
            elevation, azimuth]. All labels are given as integers. Use the
            categories, elevation_degrees, and azimuth_degrees arrays to map
            from these integers to actual values.
        """

        assert which_set in ['train', 'test']

        self.which_set = which_set

        subtensor = slice(0, stop) if stop is not None else None

        X = SmallNORB.load(which_set, 'dat', subtensor=subtensor)

        # Casts to the GPU-supported float type, using theano._asarray(), a
        # safer alternative to numpy.asarray().
        #
        # TODO: move the dtype-casting to the view_converter's output space,
        #       once dtypes-for-spaces is merged into master.
        X = theano._asarray(X, theano.config.floatX)

        # Formats data as rows in a matrix, for DenseDesignMatrix
        X = X.reshape(-1, 2 * n.prod(self.original_image_shape))

        # This is uint8
        y = SmallNORB.load(which_set, 'cat', subtensor=subtensor)
        if multi_target:
            y_extra = SmallNORB.load(which_set, 'info', subtensor=subtensor)
            y = numpy.hstack((y[:, numpy.newaxis], y_extra))

        datum_shape = ((2, ) +  # two stereo images
                       self.original_image_shape +
                       (1, ))  # one color channel

        # 's' is the stereo channel: 0 (left) or 1 (right)
        axes = ('b', 's', 0, 1, 'c')
        view_converter = StereoViewConverter(datum_shape, axes)

        super(SmallNORB, self).__init__(X=X,
                                        y=y,
                                        y_labels=numpy.max(y) + 1,
                                        view_converter=view_converter)
//  [python]
def __init__(self, which_set, multi_target=False, stop=None):
        """
        parameters
        ----------

        which_set : str
            Must be 'train' or 'test'.

        multi_target : bool
            If False, each label is an integer labeling the image catergory. If
            True, each label is a vector: [category, instance, lighting,
            elevation, azimuth]. All labels are given as integers. Use the
            categories, elevation_degrees, and azimuth_degrees arrays to map
            from these integers to actual values.
        """

        assert which_set in ['train', 'test']

        self.which_set = which_set

        subtensor = slice(0, stop) if stop is not None else None

        X = SmallNORB.load(which_set, 'dat', subtensor=subtensor)

        # Casts to the GPU-supported float type, using theano._asarray(), a
        # safer alternative to numpy.asarray().
        #
        # TODO: move the dtype-casting to the view_converter's output space,
        #       once dtypes-for-spaces is merged into master.
        X = theano._asarray(X, theano.config.floatX)

        # Formats data as rows in a matrix, for DenseDesignMatrix
        X = X.reshape(-1, 2 * numpy.prod(self.original_image_shape))

        # This is uint8
        y = SmallNORB.load(which_set, 'cat', subtensor=subtensor)
        if multi_target:
            y_extra = SmallNORB.load(which_set, 'info', subtensor=subtensor)
            y = numpy.hstack((y[:, numpy.newaxis], y_extra))

        datum_shape = ((2, ) +  # two stereo images
                       self.original_image_shape +
                       (1, ))  # one color channel

        # 's' is the stereo channel: 0 (left) or 1 (right)
        axes = ('b', 's', 0, 1, 'c')
        view_converter = StereoViewConverter(datum_shape, axes)

        super(SmallNORB, self).__init__(X=X,
                                        y=y,
                                        y_labels=numpy.max(y) + 1,
                                        view_converter=view_converter)
//  [python]
def get_output_buffers(output_shapes, device, is_float16=False):
        """Returns a dictionary of output name as key, and 1D tensor as value. The tensor has enough space for given shape."""
        data_type = torch.float16 if is_float16 else torch.float32

        output_buffers = {}
        for name, shape in output_shapes.items():
            if name == "output_selected_indices" or name == "current_step_results" or name == "last_state":
                output_buffers[name] = torch.empty(n.prod(shape), dtype=torch.long, device=device)
            elif name == "output_unfinished_sents":
                output_buffers[name] = torch.empty(n.prod(shape), dtype=torch.bool, device=device)
            else:
                output_buffers[name] = torch.empty(n.prod(shape), dtype=data_type, device=device)
        return output_buffers
//  [python]
def get_output_buffers(output_shapes, device, is_float16=False):
        """Returns a dictionary of output name as key, and 1D tensor as value. The tensor has enough space for given shape."""
        data_type = torch.float16 if is_float16 else torch.float32

        output_buffers = {}
        for name, shape in output_shapes.items():
            if name == "output_selected_indices" or name == "current_step_results" or name == "last_state":
                output_buffers[name] = torch.empty(numpy.prod(shape), dtype=torch.long, device=device)
            elif name == "output_unfinished_sents":
                output_buffers[name] = torch.empty(numpy.prod(shape), dtype=torch.bool, device=device)
            else:
                output_buffers[name] = torch.empty(numpy.prod(shape), dtype=data_type, device=device)
        return output_buffers
//  [python]
def gaussian_with_variable_parameters(
    field_dim, mean_value, sdev, mean_corr_length, std_corr_length
):
    """
    Setup a Gaussian field with correlation length drawn from a normal distribution.
    The horizontal layers are generated independently.

    Parameters:
    - field_dim        : Tuple. Dimension of the field.
    - mean_value       : Numpy array. The mean value of the field.
    - sdev             : Numpy array or float. Standard deviation of the field.
    - mean_corr_length : Float. Mean correlation length.
    - std_corr_length  : Float. Standard deviation of the correlation length.

    Returns:
    - x                : Numpy array. The generated field.
    - corr_length      : Numpy array or float. The drawn correlation length.
    """

    corr_length = add_gnoise(mean_corr_length, std_corr_length, 1)

    if len(field_dim) < 3:
        x = mean_value + fast_gaussian(field_dim, sdev, corr_length)
    else:
        layer_dim = n.prod(field_dim[:2])
        x = np.copy(mean_value)

        if np.isscalar(sdev):
            for i in range(field_dim[2]):
                idx_range = slice(i * layer_dim, (i + 1) * layer_dim)
                x[idx_range] = mean_value[idx_range] + fast_gaussian(
                    field_dim[:2], sdev, corr_length
                )
                # Generate new correlation length for the next layer
                corr_length = add_gnoise(mean_corr_length, std_corr_length, 1)
        else:
            for i in range(field_dim[2]):
                idx_range = slice(i * layer_dim, (i + 1) * layer_dim)
                x[idx_range] = mean_value[idx_range] + fast_gaussian(
                    field_dim[:2], sdev[idx_range], corr_length
                )
                # Generate new correlation length for the next layer
                corr_length = add_gnoise(mean_corr_length, std_corr_length, 1)

    return x, corr_length
//  [python]
def gaussian_with_variable_parameters(
    field_dim, mean_value, sdev, mean_corr_length, std_corr_length
):
    """
    Setup a Gaussian field with correlation length drawn from a normal distribution.
    The horizontal layers are generated independently.

    Parameters:
    - field_dim        : Tuple. Dimension of the field.
    - mean_value       : Numpy array. The mean value of the field.
    - sdev             : Numpy array or float. Standard deviation of the field.
    - mean_corr_length : Float. Mean correlation length.
    - std_corr_length  : Float. Standard deviation of the correlation length.

    Returns:
    - x                : Numpy array. The generated field.
    - corr_length      : Numpy array or float. The drawn correlation length.
    """

    corr_length = add_gnoise(mean_corr_length, std_corr_length, 1)

    if len(field_dim) < 3:
        x = mean_value + fast_gaussian(field_dim, sdev, corr_length)
    else:
        layer_dim = np.prod(field_dim[:2])
        x = np.copy(mean_value)

        if np.isscalar(sdev):
            for i in range(field_dim[2]):
                idx_range = slice(i * layer_dim, (i + 1) * layer_dim)
                x[idx_range] = mean_value[idx_range] + fast_gaussian(
                    field_dim[:2], sdev, corr_length
                )
                # Generate new correlation length for the next layer
                corr_length = add_gnoise(mean_corr_length, std_corr_length, 1)
        else:
            for i in range(field_dim[2]):
                idx_range = slice(i * layer_dim, (i + 1) * layer_dim)
                x[idx_range] = mean_value[idx_range] + fast_gaussian(
                    field_dim[:2], sdev[idx_range], corr_length
                )
                # Generate new correlation length for the next layer
                corr_length = add_gnoise(mean_corr_length, std_corr_length, 1)

    return x, corr_length
//  [python]
def forward_expected(self, inputs):
        if len(inputs) == 3:
            x, w, b = inputs
        else:
            (x, w), b = inputs, None

        n_batch_axes = self.n_batch_axes
        x_shape = self.x_shape
        w_shape = self.w_shape
        out_dtype = self.out_dtype

        if n_batch_axes is Unspecified:
            n_batch_axes = 1
        y_shape = x_shape[:n_batch_axes] + (w_shape[0],)
        x_ = x.reshape(numpy.prod(x_shape[:n_batch_axes]),
                       numpy.prod(x_shape[n_batch_axes:]))
        x_ = x_.astype(out_dtype)
        w_ = w.astype(out_dtype)
        y = x_.dot(w_.T).reshape(y_shape)
        if b is not None:
            y += b

        assert y.dtype == out_dtype
        return y,
//  [python]
def setUp_configure(self):
        from scipy import stats
        self.dist = distributions.MultivariateNormal
        self.scipy_dist = stats.multivariate_normal
        self.scipy_onebyone = True
        d, = self.event_shape

        self.test_targets = set([
            'batch_shape', 'entropy', 'event_shape', 'log_prob',
            'support'])

        loc = numpy.random.uniform(
            -1, 1, self.shape + self.event_shape).astype(numpy.float32)
        cov = numpy.random.normal(
            size=(numpy.prod(self.shape),) + (d, d))
        cov = [cov_.dot(cov_.T) for cov_ in cov]
        cov = numpy.vstack(cov).reshape(self.shape + (d, d))
        scale_tril = numpy.linalg.cholesky(cov).astype(numpy.float32)
        self.params = {'loc': loc, 'scale_tril': scale_tril}
        self.scipy_params = {'mean': loc, 'cov': cov}
//  [python]
def griewank(self, x):
        # was in [-600 600]
        x = (600./5) * x
        return 1 - n.prod(np.cos(x/sqrt(1.+np.arange(len(x))))) + sum(x**2)/4e3
//  [python]
def griewank(self, x):
        # was in [-600 600]
        x = (600./5) * x
        return 1 - np.prod(np.cos(x/sqrt(1.+np.arange(len(x))))) + sum(x**2)/4e3
//  [python]
def connectivity(grid):
    cellcenters = grid.position
    shape = numpy.asarray(cellcenters.shape)
    dim=len(shape)-1
    shape[-1] = 2**dim
    shape_with_boundary = numpy.asarray(cellcenters.shape) + 1
    shape_with_boundary = shape_with_boundary[:dim]
    indices = numpy.arange(0, numpy.multiply.reduce(shape_with_boundary), dtype=numpy.int32).reshape(shape_with_boundary)
    result = numpy.zeros(shape, dtype=numpy.int32)

    for i in range(2**dim):
        slicing1=()
        slicing2=()
        for j in range(dim):
            if i & 2**j:
                slicing2+=(slice(1,None),) 
                if len(slicing1) == 0 or slicing1[-1] is not Ellipsis:
                    slicing1+=(Ellipsis,)
            else:
                slicing2+=(slice(None,-1),)
                if len(slicing1) == 0 or slicing1[-1] is not Ellipsis:
                    slicing1+=(Ellipsis,)

        slicing1+=(i,)

        result[slicing1]=indices[slicing2]    

    return result
//  [python]
def connectivity(grid):
    cellcenters = grid.position
    shape = numpy.asarray(cellcenters.shape)
    dim=len(shape)-1
    shape[-1] = 2**dim
    shape_with_boundary = numpy.asarray(cellcenters.shape) + 1
    shape_with_boundary = shape_with_boundary[:dim]
    indices = numpy.arange(0, numpy.prod(shape_with_boundary), dtype=numpy.int32).reshape(shape_with_boundary)
    result = numpy.zeros(shape, dtype=numpy.int32)

    for i in range(2**dim):
        slicing1=()
        slicing2=()
        for j in range(dim):
            if i & 2**j:
                slicing2+=(slice(1,None),) 
                if len(slicing1) == 0 or slicing1[-1] is not Ellipsis:
                    slicing1+=(Ellipsis,)
            else:
                slicing2+=(slice(None,-1),)
                if len(slicing1) == 0 or slicing1[-1] is not Ellipsis:
                    slicing1+=(Ellipsis,)

        slicing1+=(i,)

        result[slicing1]=indices[slicing2]    

    return result
//  [python]
def sample_multinomial(N, p, size=None):
    r"""
    Draws fixed number of samples N from different
    multinomial distributions (with the same number dice sides).

    :param int N: How many samples to draw from each distribution.
    :param np.ndarray p: Probabilities specifying each distribution.
        Sum along axis 0 should be 1.
    :param size: Output shape. ``int`` or tuple of
        ``int``s. If the given shape is,
        e.g., ``(m, n, k)``, then m * n * k samples are drawn
        for each distribution.
        Default is None, in which case a single value
        is returned for each distribution.

    :rtype: np.ndarray
    :return: Array of shape ``(p.shape, size)`` or p.shape if
        size is ``None``.
    """
    # ensure s is array
    s = np.array([1]) if size is None else np.array([size]).flatten()

    def take_samples(ps):
        # we have to flatten to make apply_along_axis work.
        return np.random.multinomial(N, ps, reduce(lambda x, y: x*y, s)).flatten()

    # should have shape (prod(size)*ps.shape[0], ps.shape[1:])
    samples = np.apply_along_axis(take_samples, 0, p)
    # should have shape (size, p.shape)
    samples = samples.reshape(np.concatenate([s, p.shape]))
    # should have shape (p.shape, size)
    samples = samples.transpose(np.concatenate(
        [np.arange(s.ndim, p.ndim+s.ndim), np.arange(s.ndim)]
    ))

    if size is None:
        # get rid of trailing singleton dimension.
        samples = samples[...,0]

    return samples
//  [python]
def sample_multinomial(N, p, size=None):
    r"""
    Draws fixed number of samples N from different
    multinomial distributions (with the same number dice sides).

    :param int N: How many samples to draw from each distribution.
    :param np.ndarray p: Probabilities specifying each distribution.
        Sum along axis 0 should be 1.
    :param size: Output shape. ``int`` or tuple of
        ``int``s. If the given shape is,
        e.g., ``(m, n, k)``, then m * n * k samples are drawn
        for each distribution.
        Default is None, in which case a single value
        is returned for each distribution.

    :rtype: np.ndarray
    :return: Array of shape ``(p.shape, size)`` or p.shape if
        size is ``None``.
    """
    # ensure s is array
    s = np.array([1]) if size is None else np.array([size]).flatten()

    def take_samples(ps):
        # we have to flatten to make apply_along_axis work.
        return np.random.multinomial(N, ps, np.prod(s)).flatten()

    # should have shape (prod(size)*ps.shape[0], ps.shape[1:])
    samples = np.apply_along_axis(take_samples, 0, p)
    # should have shape (size, p.shape)
    samples = samples.reshape(np.concatenate([s, p.shape]))
    # should have shape (p.shape, size)
    samples = samples.transpose(np.concatenate(
        [np.arange(s.ndim, p.ndim+s.ndim), np.arange(s.ndim)]
    ))

    if size is None:
        # get rid of trailing singleton dimension.
        samples = samples[...,0]

    return samples
//  [python]
def check_log_prob(self, is_gpu):
        smp = self.sample_for_test()
        if is_gpu:
            log_prob1 = self.gpu_dist.log_prob(cuda.to_gpu(smp)).data
        else:
            log_prob1 = self.cpu_dist.log_prob(smp).data

        scipy_prob = self.scipy_dist.logpmf

        onebyone_smp = smp.reshape(
            (int(n.prod(self.sample_shape)),
             n.prod(self.shape),
             int(n.prod(self.event_shape))))
        onebyone_smp = n.swapaxes(onebyone_smp, 0, 1)
        onebyone_smp = onebyone_smp.reshape((-1,) + self.sample_shape
                                            + self.event_shape)

        log_prob2 = []
        for one_params, one_smp in zip(
                self.scipy_onebyone_params_iter(), onebyone_smp):
            one_smp = n.eye(3)[one_smp]
            log_prob2.append(scipy_prob(one_smp, **one_params))
        log_prob2 = n.vstack(log_prob2)
        log_prob2 = log_prob2.reshape(n.prod(self.shape), -1).T
        log_prob2 = log_prob2.reshape(self.sample_shape + self.shape)
        array.assert_allclose(log_prob1, log_prob2)
//  [python]
def check_log_prob(self, is_gpu):
        smp = self.sample_for_test()
        if is_gpu:
            log_prob1 = self.gpu_dist.log_prob(cuda.to_gpu(smp)).data
        else:
            log_prob1 = self.cpu_dist.log_prob(smp).data

        scipy_prob = self.scipy_dist.logpmf

        onebyone_smp = smp.reshape(
            (int(numpy.prod(self.sample_shape)),
             numpy.prod(self.shape),
             int(numpy.prod(self.event_shape))))
        onebyone_smp = numpy.swapaxes(onebyone_smp, 0, 1)
        onebyone_smp = onebyone_smp.reshape((-1,) + self.sample_shape
                                            + self.event_shape)

        log_prob2 = []
        for one_params, one_smp in zip(
                self.scipy_onebyone_params_iter(), onebyone_smp):
            one_smp = numpy.eye(3)[one_smp]
            log_prob2.append(scipy_prob(one_smp, **one_params))
        log_prob2 = numpy.vstack(log_prob2)
        log_prob2 = log_prob2.reshape(numpy.prod(self.shape), -1).T
        log_prob2 = log_prob2.reshape(self.sample_shape + self.shape)
        array.assert_allclose(log_prob1, log_prob2)
//  [python]
def check_sample(self, is_gpu):
        if is_gpu:
            smp1 = self.gpu_dist.sample(
                sample_shape=(100000,)+self.sample_shape).data
        else:
            smp1 = self.cpu_dist.sample(
                sample_shape=(100000,)+self.sample_shape).data

        smp2 = []
        for one_params in self.scipy_onebyone_params_iter():
            smp2.append(self.scipy_dist.rvs(
                size=(100000,)+self.sample_shape, **one_params))
        smp2 = numpy.vstack(smp2)
        smp2 = smp2.dot(numpy.arange(3))
        smp2 = smp2.reshape((numpy.prod(self.shape), 100000)
                            + self.sample_shape
                            + self.cpu_dist.event_shape)
        smp2 = numpy.rollaxis(
            smp2, 0, smp2.ndim-len(self.cpu_dist.event_shape))
        smp2 = smp2.reshape((100000,) + self.sample_shape + self.shape
                            + self.cpu_dist.event_shape)
        array.assert_allclose(smp1.mean(axis=0), smp2.mean(axis=0),
                              atol=3e-2, rtol=3e-2)
        array.assert_allclose(smp1.std(axis=0), smp2.std(axis=0),
                              atol=3e-2, rtol=3e-2)
//  [python]
def loglikelihood(self, x, previous=False):
        """return log-likelihood of `x` regarding the current sample distribution"""
        # testing of original fct: MC integrate must be one: mean(p(x_i)) * volume(where x_i are uniformely sampled)
        # for i in range(3): print mean([cma.likelihood(20*r-10, dim * [0], None, 3) for r in rand(10000,dim)]) * 20**dim
        # TODO: test this!!
        # c=cma.fmin...
        # c[3]['cma'].loglikelihood(...)

        if previous and hasattr(self, 'lastiter'):
            sigma = self.lastiter.sigma
            Crootinv = self.lastiter._Crootinv
            xmean = self.lastiter.mean
            D = self.lastiter.D
        elif previous and self.countiter > 1:
            raise _Error('no previous distribution parameters stored, check options importance_mixing')
        else:
            sigma = self.sigma
            Crootinv = self._Crootinv
            xmean = self.mean
            D = self.D

        dx = array(x) - xmean  # array(x) - array(m)
        n = self.N
        logs2pi = n * log(2*np.pi) / 2.
        logdetC = 2 * sum(log(D))
        dx = np.dot(Crootinv, dx)
        res = -sum(dx**2) / sigma**2 / 2 - logs2pi - logdetC/2 - n*log(sigma)
        if 1 < 3: # testing
            s2pi = (2*np.pi)**(n/2.)
            detC = n.prod(D)**2
            res2 = -sum(dx**2) / sigma**2 / 2 - log(s2pi * abs(detC)**0.5 * sigma**n)
            assert res2 < res + 1e-8 or res2 > res - 1e-8
        return res
//  [python]
def loglikelihood(self, x, previous=False):
        """return log-likelihood of `x` regarding the current sample distribution"""
        # testing of original fct: MC integrate must be one: mean(p(x_i)) * volume(where x_i are uniformely sampled)
        # for i in range(3): print mean([cma.likelihood(20*r-10, dim * [0], None, 3) for r in rand(10000,dim)]) * 20**dim
        # TODO: test this!!
        # c=cma.fmin...
        # c[3]['cma'].loglikelihood(...)

        if previous and hasattr(self, 'lastiter'):
            sigma = self.lastiter.sigma
            Crootinv = self.lastiter._Crootinv
            xmean = self.lastiter.mean
            D = self.lastiter.D
        elif previous and self.countiter > 1:
            raise _Error('no previous distribution parameters stored, check options importance_mixing')
        else:
            sigma = self.sigma
            Crootinv = self._Crootinv
            xmean = self.mean
            D = self.D

        dx = array(x) - xmean  # array(x) - array(m)
        n = self.N
        logs2pi = n * log(2*np.pi) / 2.
        logdetC = 2 * sum(log(D))
        dx = np.dot(Crootinv, dx)
        res = -sum(dx**2) / sigma**2 / 2 - logs2pi - logdetC/2 - n*log(sigma)
        if 1 < 3: # testing
            s2pi = (2*np.pi)**(n/2.)
            detC = np.prod(D)**2
            res2 = -sum(dx**2) / sigma**2 / 2 - log(s2pi * abs(detC)**0.5 * sigma**n)
            assert res2 < res + 1e-8 or res2 > res - 1e-8
        return res
//  [python]
def offset_normal_chains_should_fail_wrapper(
      self, sample_shape, independent_chain_shape, other_shape):
    """Check results with normal chains that are offset from each other."""

    state_shape = sample_shape + independent_chain_shape + other_shape
    state_ = np.random.randn(*state_shape)

    # Add a significant offset to the different (formerly iid) chains.
    offset = np.linspace(
        0, 2, num=np.prod(independent_chain_shape)).reshape([1] * len(
            sample_shape) + independent_chain_shape + [1] * len(other_shape))
    state_ += offset

    self.check_results(state_, independent_chain_shape, should_pass=False)
//  [python]
def __init__(self, *args):
        supported_formats = (np.int8, np.int16, np.int32, np.int64)

        # Convert input to numpy
        array_list = [np.asarray(a.cpu()) if isinstance(a, torch.Tensor)
                      else np.asarray(a)
                      for a in args]
        assert array_list[0].ndim == 1, \
            'Only 1D arrays are accepted as input.'
        assert all([a.shape == array_list[0].shape for a in array_list]), \
            'All input arrays must have the same shape'
        assert all([a.dtype in supported_formats for a in array_list]), \
            f'All input arrays must be in {supported_formats}. ' \
            f'Received types: {[a.dtype in supported_formats for a in array_list]}'

        # Compute the bases to build the composite array
        dtype_list = [a.dtype for a in array_list]
        dtype_max_list = [np.iinfo(dt).max for dt in dtype_list]
        dtype_max = max(dtype_max_list)
        dtype = dtype_list[dtype_max_list.index(dtype_max)]
        if array_list[0].shape[0] == 0:
            max_list = [0 for _ in array_list]
        else:
            max_list = [np.abs(a).max() + 1 for a in array_list]
        assert all([reduce(lambda x, y: x * y, max_list) < dtype_max]), \
            'The dtype of at least one of the input arrays must ' \
            'allow the composite computation.'
        base_list = [reduce(lambda x, y: x * y, max_list[i + 1:])
                     for i in range(len(array_list) - 1)] + [1]

        # Build the composite array
        self.dtype_list = dtype_list
        self.dtype = dtype
        self.max_list = max_list
        self.base_list = base_list
        self.data = sum([a.astype(dtype) * b
                         for a, b in zip(array_list, base_list)])
//  [python]
def __init__(self, *args):
        supported_formats = (np.int8, np.int16, np.int32, np.int64)

        # Convert input to numpy
        array_list = [np.asarray(a.cpu()) if isinstance(a, torch.Tensor)
                      else np.asarray(a)
                      for a in args]
        assert array_list[0].ndim == 1, \
            'Only 1D arrays are accepted as input.'
        assert all([a.shape == array_list[0].shape for a in array_list]), \
            'All input arrays must have the same shape'
        assert all([a.dtype in supported_formats for a in array_list]), \
            f'All input arrays must be in {supported_formats}. ' \
            f'Received types: {[a.dtype in supported_formats for a in array_list]}'

        # Compute the bases to build the composite array
        dtype_list = [a.dtype for a in array_list]
        dtype_max_list = [np.iinfo(dt).max for dt in dtype_list]
        dtype_max = max(dtype_max_list)
        dtype = dtype_list[dtype_max_list.index(dtype_max)]
        if array_list[0].shape[0] == 0:
            max_list = [0 for _ in array_list]
        else:
            max_list = [np.abs(a).max() + 1 for a in array_list]
        assert all([np.prod(max_list) < dtype_max]), \
            'The dtype of at least one of the input arrays must ' \
            'allow the composite computation.'
        base_list = [np.prod(max_list[i + 1:])
                     for i in range(len(array_list) - 1)] + [1]

        # Build the composite array
        self.dtype_list = dtype_list
        self.dtype = dtype
        self.max_list = max_list
        self.base_list = base_list
        self.data = sum([a.astype(dtype) * b
                         for a, b in zip(array_list, base_list)])
//  [python]
def batched_gesv(a, b):
    """Solves multiple linear matrix equations using cublas<t>getr[fs]Batched().

    Computes the solution to system of linear equation ``ax = b``.

    Args:
        a (cupy.ndarray): The matrix with dimension ``(..., M, M)``.
        b (cupy.ndarray): The matrix with dimension ``(..., M)`` or
            ``(..., M, K)``.

    Returns:
        cupy.ndarray:
            The matrix with dimension ``(..., M)`` or ``(..., M, K)``.
    """  # NOQA
    _util._assert_cupy_array(a, b)
    _util._assert_stacked_2d(a)
    _util._assert_stacked_square(a)

    # TODO(kataoka): Support broadcast
    if not (
        (a.ndim == b.ndim or a.ndim == b.ndim + 1)
        and a.shape[:-1] == b.shape[:a.ndim - 1]
    ):
        raise ValueError(
            'a must have (..., M, M) shape and b must have (..., M) '
            'or (..., M, K)')

    dtype, out_dtype = _util.linalg_common_type(a, b)
    if b.size == 0:
        return cupy.empty(b.shape, out_dtype)

    if dtype == 'f':
        t = 's'
    elif dtype == 'd':
        t = 'd'
    elif dtype == 'F':
        t = 'c'
    elif dtype == 'D':
        t = 'z'
    else:
        raise TypeError('invalid dtype')
    getrf = getattr(cublas, t + 'getrfBatched')
    getrs = getattr(cublas, t + 'getrsBatched')

    bs = n if a.ndim > 2 else 1
    n = a.shape[-1]
    nrhs = b.shape[-1] if a.ndim == b.ndim else 1
    b_shape = b.shape
    a_data_ptr = a.data.ptr
    b_data_ptr = b.data.ptr
    a = cupy.ascontiguousarray(a.reshape(bs, n, n).transpose(0, 2, 1),
                               dtype=dtype)
    b = cupy.ascontiguousarray(b.reshape(bs, n, nrhs).transpose(0, 2, 1),
                               dtype=dtype)
    if a.data.ptr == a_data_ptr:
        a = a.copy()
    if b.data.ptr == b_data_ptr:
        b = b.copy()

    if n > get_batched_gesv_limit():
        warnings.warn('The matrix size ({}) exceeds the set limit ({})'.
                      format(n, get_batched_gesv_limit()))

    handle = device.get_cublas_handle()
    lda = n
    a_step = lda * n * a.itemsize
    a_array = cupy.arange(a.data.ptr, a.data.ptr + a_step * bs, a_step,
                          dtype=cupy.uintp)
    ldb = n
    b_step = ldb * nrhs * b.itemsize
    b_array = cupy.arange(b.data.ptr, b.data.ptr + b_step * bs, b_step,
                          dtype=cupy.uintp)
    pivot = cupy.empty((bs, n), dtype=numpy.int32)
    dinfo = cupy.empty((bs,), dtype=numpy.int32)
    info = numpy.empty((1,), dtype=numpy.int32)
    # LU factorization (A = L * U)
    getrf(handle, n, a_array.data.ptr, lda, pivot.data.ptr, dinfo.data.ptr, bs)
    _util._check_cublas_info_array_if_synchronization_allowed(getrf, dinfo)
    # Solves Ax = b
    getrs(handle, cublas.CUBLAS_OP_N, n, nrhs, a_array.data.ptr, lda,
          pivot.data.ptr, b_array.data.ptr, ldb, info.ctypes.data, bs)
    if info[0] != 0:
        msg = 'Error reported by {} in cuBLAS. '.format(getrs.__name__)
        if info[0] < 0:
            msg += 'The {}-th parameter had an illegal value.'.format(-info[0])
        raise linalg.LinAlgError(msg)

    return b.transpose(0, 2, 1).reshape(b_shape).astype(out_dtype, copy=False)
//  [python]
def batched_gesv(a, b):
    """Solves multiple linear matrix equations using cublas<t>getr[fs]Batched().

    Computes the solution to system of linear equation ``ax = b``.

    Args:
        a (cupy.ndarray): The matrix with dimension ``(..., M, M)``.
        b (cupy.ndarray): The matrix with dimension ``(..., M)`` or
            ``(..., M, K)``.

    Returns:
        cupy.ndarray:
            The matrix with dimension ``(..., M)`` or ``(..., M, K)``.
    """  # NOQA
    _util._assert_cupy_array(a, b)
    _util._assert_stacked_2d(a)
    _util._assert_stacked_square(a)

    # TODO(kataoka): Support broadcast
    if not (
        (a.ndim == b.ndim or a.ndim == b.ndim + 1)
        and a.shape[:-1] == b.shape[:a.ndim - 1]
    ):
        raise ValueError(
            'a must have (..., M, M) shape and b must have (..., M) '
            'or (..., M, K)')

    dtype, out_dtype = _util.linalg_common_type(a, b)
    if b.size == 0:
        return cupy.empty(b.shape, out_dtype)

    if dtype == 'f':
        t = 's'
    elif dtype == 'd':
        t = 'd'
    elif dtype == 'F':
        t = 'c'
    elif dtype == 'D':
        t = 'z'
    else:
        raise TypeError('invalid dtype')
    getrf = getattr(cublas, t + 'getrfBatched')
    getrs = getattr(cublas, t + 'getrsBatched')

    bs = numpy.prod(a.shape[:-2]) if a.ndim > 2 else 1
    n = a.shape[-1]
    nrhs = b.shape[-1] if a.ndim == b.ndim else 1
    b_shape = b.shape
    a_data_ptr = a.data.ptr
    b_data_ptr = b.data.ptr
    a = cupy.ascontiguousarray(a.reshape(bs, n, n).transpose(0, 2, 1),
                               dtype=dtype)
    b = cupy.ascontiguousarray(b.reshape(bs, n, nrhs).transpose(0, 2, 1),
                               dtype=dtype)
    if a.data.ptr == a_data_ptr:
        a = a.copy()
    if b.data.ptr == b_data_ptr:
        b = b.copy()

    if n > get_batched_gesv_limit():
        warnings.warn('The matrix size ({}) exceeds the set limit ({})'.
                      format(n, get_batched_gesv_limit()))

    handle = device.get_cublas_handle()
    lda = n
    a_step = lda * n * a.itemsize
    a_array = cupy.arange(a.data.ptr, a.data.ptr + a_step * bs, a_step,
                          dtype=cupy.uintp)
    ldb = n
    b_step = ldb * nrhs * b.itemsize
    b_array = cupy.arange(b.data.ptr, b.data.ptr + b_step * bs, b_step,
                          dtype=cupy.uintp)
    pivot = cupy.empty((bs, n), dtype=numpy.int32)
    dinfo = cupy.empty((bs,), dtype=numpy.int32)
    info = numpy.empty((1,), dtype=numpy.int32)
    # LU factorization (A = L * U)
    getrf(handle, n, a_array.data.ptr, lda, pivot.data.ptr, dinfo.data.ptr, bs)
    _util._check_cublas_info_array_if_synchronization_allowed(getrf, dinfo)
    # Solves Ax = b
    getrs(handle, cublas.CUBLAS_OP_N, n, nrhs, a_array.data.ptr, lda,
          pivot.data.ptr, b_array.data.ptr, ldb, info.ctypes.data, bs)
    if info[0] != 0:
        msg = 'Error reported by {} in cuBLAS. '.format(getrs.__name__)
        if info[0] < 0:
            msg += 'The {}-th parameter had an illegal value.'.format(-info[0])
        raise linalg.LinAlgError(msg)

    return b.transpose(0, 2, 1).reshape(b_shape).astype(out_dtype, copy=False)