jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _write_kw(fileobj, propname, values, lookup, dimens):
    """Write a section for a single keyword to an already open file.

    :param fileobj:   Where to serialize the data
    :param propname:  Name of the property from the grid to write
    :param values:    Array with the values themselves
    :param lookup:    Mapping of properties to identfier string used in file
    :param dimens:    Dimensions of the grid, (ni, nj, nk)
    """
    # structure of the values: if we can make a fixed number of columns,
    # then we attempt to do so, otherwise we just write everything as a
    # long stream (faster to load, when there are no exceptions)
    if propname in _SPECIAL_SHAPE:
        shape = _SPECIAL_SHAPE[propname](*dimens)
    else:
        shape = (numpy.multiply.reduce(values.shape), 1)

    # make sure we get an appropriate number of digits on the keywords that
    # are known to us
    if propname in _SPECIAL_FORMAT:
        fmt = _SPECIAL_FORMAT[propname]
    else:
        # print integral types without inappropriate fraction digits
        if numpy.issubdtype(values.dtype, numpy.integer):
            fmt = '%d'
        else:
            fmt = '%11.6f'

    # write the keyword itself, on its own line
    fileobj.write(enc('{0:<8s}\n'.format(lookup(propname))))

    # replace NaN values (which we typically use to trap access to inactive
    # blocks) with zero (which Petrel uses a placeholder for inactive blocks)
    zero = 0 if numpy.issubdtype(values.dtype, numpy.integer) else 0.
    data = numpy.reshape(numpy.copy(numpy.ma.getdata(values)), shape)
    data[numpy.isnan(data)] = zero

    # data array itself; NumPy does the job for us
    numpy.savetxt(fileobj, data, fmt=fmt, delimiter='\t')

    # end of field for most fields in the loaded grid file
    fileobj.write(enc('{0:s}\n'.format(lookup('/'))))
//  [python]
def _write_kw(fileobj, propname, values, lookup, dimens):
    """Write a section for a single keyword to an already open file.

    :param fileobj:   Where to serialize the data
    :param propname:  Name of the property from the grid to write
    :param values:    Array with the values themselves
    :param lookup:    Mapping of properties to identfier string used in file
    :param dimens:    Dimensions of the grid, (ni, nj, nk)
    """
    # structure of the values: if we can make a fixed number of columns,
    # then we attempt to do so, otherwise we just write everything as a
    # long stream (faster to load, when there are no exceptions)
    if propname in _SPECIAL_SHAPE:
        shape = _SPECIAL_SHAPE[propname](*dimens)
    else:
        shape = (numpy.prod(values.shape), 1)

    # make sure we get an appropriate number of digits on the keywords that
    # are known to us
    if propname in _SPECIAL_FORMAT:
        fmt = _SPECIAL_FORMAT[propname]
    else:
        # print integral types without inappropriate fraction digits
        if numpy.issubdtype(values.dtype, numpy.integer):
            fmt = '%d'
        else:
            fmt = '%11.6f'

    # write the keyword itself, on its own line
    fileobj.write(enc('{0:<8s}\n'.format(lookup(propname))))

    # replace NaN values (which we typically use to trap access to inactive
    # blocks) with zero (which Petrel uses a placeholder for inactive blocks)
    zero = 0 if numpy.issubdtype(values.dtype, numpy.integer) else 0.
    data = numpy.reshape(numpy.copy(numpy.ma.getdata(values)), shape)
    data[numpy.isnan(data)] = zero

    # data array itself; NumPy does the job for us
    numpy.savetxt(fileobj, data, fmt=fmt, delimiter='\t')

    # end of field for most fields in the loaded grid file
    fileobj.write(enc('{0:s}\n'.format(lookup('/'))))
//  [python]
def _update_shape_range(out, x, shape_v, neg_index, dim_prod):
        """ update min and max shape of output when input shape is dynamic"""
        x_min_shape = x['shape']
        x_max_shape = x['shape']
        if 'max_shape' in x:
            x_max_shape = x['max_shape']
        if 'min_shape' in x:
            x_min_shape = x['min_shape']
        max_arr_prod = n.prod(x_max_shape)
        min_arr_prod = n.prod(x_min_shape)
        max_shape = list(shape_v)
        min_shape = list(shape_v)
        if neg_index != -1:
            max_shape[neg_index] = int(max_arr_prod / dim_prod)
            min_shape[neg_index] = int(min_arr_prod / dim_prod)
        out['max_shape'] = tuple(max_shape)
        out['min_shape'] = tuple(min_shape)
        return out
//  [python]
def _update_shape_range(out, x, shape_v, neg_index, dim_prod):
        """ update min and max shape of output when input shape is dynamic"""
        x_min_shape = x['shape']
        x_max_shape = x['shape']
        if 'max_shape' in x:
            x_max_shape = x['max_shape']
        if 'min_shape' in x:
            x_min_shape = x['min_shape']
        max_arr_prod = np.prod(x_max_shape)
        min_arr_prod = np.prod(x_min_shape)
        max_shape = list(shape_v)
        min_shape = list(shape_v)
        if neg_index != -1:
            max_shape[neg_index] = int(max_arr_prod / dim_prod)
            min_shape[neg_index] = int(min_arr_prod / dim_prod)
        out['max_shape'] = tuple(max_shape)
        out['min_shape'] = tuple(min_shape)
        return out
//  [python]
def _get_shape_and_range(x, shape):
        """ get min and max shape when output shape is dynamic"""
        min_shape = None
        max_shape = None
        x_shp = x['shape']
        if is_shape_unknown(shape['shape']):
            out_shape = [-2]
            return out_shape, min_shape, max_shape

        shape_rank = shape['shape'][0]
        if not x_shp:
            # x is a scalar, output shape fixed
            out_shape = [1] * shape_rank
            return out_shape, min_shape, max_shape

        out_shape = [-1] * shape_rank
        if "max_value" in shape and "min_value" in shape:
            min_shape = shape["min_value"]
            max_shape = shape["max_value"]
            if len(min_shape) != shape_rank or len(max_shape) != shape_rank:
                raise RuntimeError("The primitive[Reshape]'s input[shape] min or max value not math the shape rank.")
            for i in range(shape_rank):
                if min_shape[i] == max_shape[i] and min_shape[i] != 1:
                    out_shape[i] = min_shape[i]
        elif is_shape_unknown(x_shp) and "max_shape" in x:
            # when dynamic memory allocation is supported, max_shape can be left out
            min_shape = [1] * shape_rank
            max_shape = [int(n.prod(x["max_shape"]))] * shape_rank
        return out_shape, min_shape, max_shape
//  [python]
def _get_shape_and_range(x, shape):
        """ get min and max shape when output shape is dynamic"""
        min_shape = None
        max_shape = None
        x_shp = x['shape']
        if is_shape_unknown(shape['shape']):
            out_shape = [-2]
            return out_shape, min_shape, max_shape

        shape_rank = shape['shape'][0]
        if not x_shp:
            # x is a scalar, output shape fixed
            out_shape = [1] * shape_rank
            return out_shape, min_shape, max_shape

        out_shape = [-1] * shape_rank
        if "max_value" in shape and "min_value" in shape:
            min_shape = shape["min_value"]
            max_shape = shape["max_value"]
            if len(min_shape) != shape_rank or len(max_shape) != shape_rank:
                raise RuntimeError("The primitive[Reshape]'s input[shape] min or max value not math the shape rank.")
            for i in range(shape_rank):
                if min_shape[i] == max_shape[i] and min_shape[i] != 1:
                    out_shape[i] = min_shape[i]
        elif is_shape_unknown(x_shp) and "max_shape" in x:
            # when dynamic memory allocation is supported, max_shape can be left out
            min_shape = [1] * shape_rank
            max_shape = [int(np.prod(x["max_shape"]))] * shape_rank
        return out_shape, min_shape, max_shape
//  [python]
def get_group_index(labels, shape, sort, xnull):
    """
    For the particular label_list, gets the offsets into the hypothetical list
    representing the totally ordered cartesian product of all possible label
    combinations, *as long as* this space fits within int64 bounds;
    otherwise, though group indices identify unique combinations of
    labels, they cannot be deconstructed.
    - If `sort`, rank of returned ids preserve lexical ranks of labels.
      i.e. returned id's can be used to do lexical sort on labels;
    - If `xnull` nulls (-1 labels) are passed through.

    Parameters
    ----------
    labels: sequence of arrays
        Integers identifying levels at each location
    shape: sequence of ints same length as labels
        Number of unique levels at each location
    sort: boolean
        If the ranks of returned ids should match lexical ranks of labels
    xnull: boolean
        If true nulls are excluded. i.e. -1 values in the labels are
        passed through
    Returns
    -------
    An array of type int64 where two elements are equal if their corresponding
    labels are equal at all location.
    """
    def _int64_cut_off(shape):
        acc = long(1)
        for i, mul in enumerate(shape):
            acc *= long(mul)
            if not acc < _INT64_MAX:
                return i
        return len(shape)

    def loop(labels, shape):
        # how many levels can be done without overflow:
        nlev = _int64_cut_off(shape)

        # compute flat ids for the first `nlev` levels
        stride = n.prod(shape[1:nlev], dtype='i8')
        out = stride * labels[0].astype('i8', subok=False, copy=False)

        for i in range(1, nlev):
            stride //= shape[i]
            out += labels[i] * stride

        if xnull: # exclude nulls
            mask = labels[0] == -1
            for lab in labels[1:nlev]:
                mask |= lab == -1
            out[mask] = -1

        if nlev == len(shape):  # all levels done!
            return out

        # compress what has been done so far in order to avoid overflow
        # to retain lexical ranks, obs_ids should be sorted
        comp_ids, obs_ids = _compress_group_index(out, sort=sort)

        labels = [comp_ids] + labels[nlev:]
        shape = [len(obs_ids)] + shape[nlev:]

        return loop(labels, shape)

    def maybe_lift(lab, size):  # pormote nan values
        return (lab + 1, size + 1) if (lab == -1).any() else (lab, size)

    labels = map(com._ensure_int64, labels)
    if not xnull:
        labels, shape = map(list, zip(*map(maybe_lift, labels, shape)))

    return loop(list(labels), list(shape))
//  [python]
def get_group_index(labels, shape, sort, xnull):
    """
    For the particular label_list, gets the offsets into the hypothetical list
    representing the totally ordered cartesian product of all possible label
    combinations, *as long as* this space fits within int64 bounds;
    otherwise, though group indices identify unique combinations of
    labels, they cannot be deconstructed.
    - If `sort`, rank of returned ids preserve lexical ranks of labels.
      i.e. returned id's can be used to do lexical sort on labels;
    - If `xnull` nulls (-1 labels) are passed through.

    Parameters
    ----------
    labels: sequence of arrays
        Integers identifying levels at each location
    shape: sequence of ints same length as labels
        Number of unique levels at each location
    sort: boolean
        If the ranks of returned ids should match lexical ranks of labels
    xnull: boolean
        If true nulls are excluded. i.e. -1 values in the labels are
        passed through
    Returns
    -------
    An array of type int64 where two elements are equal if their corresponding
    labels are equal at all location.
    """
    def _int64_cut_off(shape):
        acc = long(1)
        for i, mul in enumerate(shape):
            acc *= long(mul)
            if not acc < _INT64_MAX:
                return i
        return len(shape)

    def loop(labels, shape):
        # how many levels can be done without overflow:
        nlev = _int64_cut_off(shape)

        # compute flat ids for the first `nlev` levels
        stride = np.prod(shape[1:nlev], dtype='i8')
        out = stride * labels[0].astype('i8', subok=False, copy=False)

        for i in range(1, nlev):
            stride //= shape[i]
            out += labels[i] * stride

        if xnull: # exclude nulls
            mask = labels[0] == -1
            for lab in labels[1:nlev]:
                mask |= lab == -1
            out[mask] = -1

        if nlev == len(shape):  # all levels done!
            return out

        # compress what has been done so far in order to avoid overflow
        # to retain lexical ranks, obs_ids should be sorted
        comp_ids, obs_ids = _compress_group_index(out, sort=sort)

        labels = [comp_ids] + labels[nlev:]
        shape = [len(obs_ids)] + shape[nlev:]

        return loop(labels, shape)

    def maybe_lift(lab, size):  # pormote nan values
        return (lab + 1, size + 1) if (lab == -1).any() else (lab, size)

    labels = map(com._ensure_int64, labels)
    if not xnull:
        labels, shape = map(list, zip(*map(maybe_lift, labels, shape)))

    return loop(list(labels), list(shape))
//  [python]
def max_pairs(shape):
    """[DEPRECATED] Compute the maximum number of record pairs possible."""

    if not isinstance(shape, (tuple, list)):
        x = get_length(shape)
        n = int(x * (x - 1) / 2)

    elif isinstance(shape, (tuple, list)) and len(shape) == 1:
        x = get_length(shape[0])
        n = int(x * (x - 1) / 2)

    else:
        n = 1
        for xi in shape:
            n *= get_length(xi)

    return n
//  [python]
def max_pairs(shape):
    """[DEPRECATED] Compute the maximum number of record pairs possible."""

    if not isinstance(shape, (tuple, list)):
        x = get_length(shape)
        n = int(x * (x - 1) / 2)

    elif isinstance(shape, (tuple, list)) and len(shape) == 1:
        x = get_length(shape[0])
        n = int(x * (x - 1) / 2)

    else:
        n = numpy.prod([get_length(xi) for xi in shape])

    return n
//  [python]
def _update_shape_and_value(self, out, x, shape_v, dim_prod, neg_index):
        """ update shape, value and min / max value of output when input shape is known"""
        x_shp = x['shape']
        if dim_prod <= 0:
            raise ValueError(f"For '{self.name}', the shape of 'input_x' is {x_shp}, "
                             f"the value of 'input_shape' is {shape_v}. "
                             f"The product of 'input_shape' should > 0, but got {dim_prod}.")
        arr_prod = n.prod(x_shp)
        if neg_index != -1:
            shape_v[neg_index] = int(arr_prod / dim_prod)
            dim_prod *= shape_v[neg_index]
        if dim_prod != arr_prod:
            raise ValueError(f"For '{self.name}', the shape of 'input_x' is {x_shp}, "
                             f"the value of 'input_shape' value is {shape_v}. "
                             f"The product of the shape of 'input_x' should be equal to that of 'input_shape', "
                             f"but got {arr_prod} and {dim_prod}.")
        out['shape'] = tuple(shape_v)

        if x['value'] is not None:
            out['value'] = Tensor(x['value'].asnumpy().reshape(shape_v))

        if ('min_value' in x and 'max_value' in x):
            ret_min_value = np.array(x['min_value']).reshape(shape_v)
            ret_max_value = np.array(x['max_value']).reshape(shape_v)
            ret_min_value = tuple(ret_min_value.tolist())
            ret_max_value = tuple(ret_max_value.tolist())
            out['min_value'] = ret_min_value
            out['max_value'] = ret_max_value
        return out
//  [python]
def _update_shape_and_value(self, out, x, shape_v, dim_prod, neg_index):
        """ update shape, value and min / max value of output when input shape is known"""
        x_shp = x['shape']
        if dim_prod <= 0:
            raise ValueError(f"For '{self.name}', the shape of 'input_x' is {x_shp}, "
                             f"the value of 'input_shape' is {shape_v}. "
                             f"The product of 'input_shape' should > 0, but got {dim_prod}.")
        arr_prod = np.prod(x_shp)
        if neg_index != -1:
            shape_v[neg_index] = int(arr_prod / dim_prod)
            dim_prod *= shape_v[neg_index]
        if dim_prod != arr_prod:
            raise ValueError(f"For '{self.name}', the shape of 'input_x' is {x_shp}, "
                             f"the value of 'input_shape' value is {shape_v}. "
                             f"The product of the shape of 'input_x' should be equal to that of 'input_shape', "
                             f"but got {arr_prod} and {dim_prod}.")
        out['shape'] = tuple(shape_v)

        if x['value'] is not None:
            out['value'] = Tensor(x['value'].asnumpy().reshape(shape_v))

        if ('min_value' in x and 'max_value' in x):
            ret_min_value = np.array(x['min_value']).reshape(shape_v)
            ret_max_value = np.array(x['max_value']).reshape(shape_v)
            ret_min_value = tuple(ret_min_value.tolist())
            ret_max_value = tuple(ret_max_value.tolist())
            out['min_value'] = ret_min_value
            out['max_value'] = ret_max_value
        return out
//  [python]
def volume(self) -> float:
        return float(np.prod(self.size))
//  [python]
def _reduce_jvp(reducer, init_values, primals, tangents, axes):
  input_shape = np.array(primals[0].shape, dtype=int)

  n = np.prod(input_shape[list(axes)])
  non_axes = np.delete(np.arange(len(input_shape)), axes)

  # Move the reduced axes to the front, and flatten them to 1D.
  permutation = axes + tuple(non_axes)
  new_shape = (n,) + tuple(input_shape[non_axes])
  primals = tuple(reshape(x, new_shape, permutation) for x in primals)
  tangents = tuple(reshape(t, new_shape, permutation) for t in tangents)

  for d in range(len(non_axes) + 1):
    reducer = api.vmap(reducer)
  def _reduce_tree(*xs, axis=0):
    """Reduce by repeatedly splitting the array and multiplying."""
    while xs[0].shape[axis] > 1:
      n = xs[0].shape[axis]
      n1 = (n + 1) // 2
      n2 = n - n1
      xs1 = [slicing.slice_in_dim(x, 0, n1) for x in xs]
      xs2 = [slicing.slice_in_dim(x, n1, None) for x in xs]
      if n2 != n1:
        paddings = [(0, 0, 0)] * len(xs[0].shape)
        paddings[axis] = (0, 1, 0)
        xs2 = [pad(x2, i, paddings) for x2, i in zip(xs2, init_values)]
      xs = reducer(*(xs1 + xs2))
    if xs[0].shape[axis] == 0:
      return [full(input_shape[non_axes], i) for i in init_values]
    return tuple(squeeze(x, (axis,)) for x in xs)

  return api.jvp(_reduce_tree, primals, tangents)
//  [python]
def full_index_size(*args):
    """Compute the number of records in a full index.

    Compute the number of records in a full index without building the index
    itself. The result is the maximum number of record pairs possible. This
    function is especially useful in measures like the `reduction_ratio`.

    Deduplication: Given a DataFrame A with length N, the full index size is
    N*(N-1)/2. Linking: Given a DataFrame A with length N and a DataFrame B
    with length M, the full index size is N*M.

    Parameters
    ----------
    *args: int, pandas.MultiIndex, pandas.Series, pandas.DataFrame
        A pandas object or a int representing the length of a dataset to link.
        When there is one argument, it is assumed that the record linkage is
        a deduplication process.

    Examples
    --------

    Use integers:
    >>> full_index_size(10)  # deduplication: 45 pairs
    >>> full_index_size(10, 10)  # linking: 100 pairs

    or pandas objects
    >>> full_index_size(DF)  # deduplication: len(DF)*(len(DF)-1)/2 pairs
    >>> full_index_size(DF, DF)  # linking: len(DF)*len(DF) pairs

    """

    # check if a list or tuple is passed as argument
    if len(args) == 1 and isinstance(args[0], (list, tuple)):
        args = tuple(args[0])

    if len(args) == 1:
        n = get_length(args[0])
        size = int(n * (n - 1) / 2)
    else:
        lengths = [get_length(arg) for arg in args]
        size = 1
        for n in lengths:
            size *= n

    return size
//  [python]
def full_index_size(*args):
    """Compute the number of records in a full index.

    Compute the number of records in a full index without building the index
    itself. The result is the maximum number of record pairs possible. This
    function is especially useful in measures like the `reduction_ratio`.

    Deduplication: Given a DataFrame A with length N, the full index size is
    N*(N-1)/2. Linking: Given a DataFrame A with length N and a DataFrame B
    with length M, the full index size is N*M.

    Parameters
    ----------
    *args: int, pandas.MultiIndex, pandas.Series, pandas.DataFrame
        A pandas object or a int representing the length of a dataset to link.
        When there is one argument, it is assumed that the record linkage is
        a deduplication process.

    Examples
    --------

    Use integers:
    >>> full_index_size(10)  # deduplication: 45 pairs
    >>> full_index_size(10, 10)  # linking: 100 pairs

    or pandas objects
    >>> full_index_size(DF)  # deduplication: len(DF)*(len(DF)-1)/2 pairs
    >>> full_index_size(DF, DF)  # linking: len(DF)*len(DF) pairs

    """

    # check if a list or tuple is passed as argument
    if len(args) == 1 and isinstance(args[0], (list, tuple)):
        args = tuple(args[0])

    if len(args) == 1:
        n = get_length(args[0])
        size = int(n * (n - 1) / 2)
    else:
        size = numpy.prod([get_length(arg) for arg in args])

    return size
//  [python]
def numpy_ndarray_resolver(obj):
    meta = obj.meta
    value_type_name = meta['value_type_']
    if value_type_name == 'object':
        view = memoryview(obj.member('buffer_'))
        return pickle.loads(view, fix_imports=True)

    value_type = normalize_dtype(value_type_name, meta.get('value_type_meta_', None))
    # process string ndarray from c++
    if value_type_name in ['str', 'string', 'std::string', 'std::__1::string']:
        from vineyard.data.arrow import binary_array_resolver
        from vineyard.data.arrow import string_array_resolver

        try:
            return string_array_resolver(obj.member('buffer_')).to_numpy(
                zero_copy_only=False
            )
        except:  # noqa: E722, pylint: disable=bare-except
            return binary_array_resolver(obj.member('buffer_')).to_numpy(
                zero_copy_only=False
            )

    shape = from_json(meta['shape_'])
    if 'order_' in meta:
        order = from_json(meta['order_'])
    else:
        order = 'C'
    if n.prod(shape) == 0:
        return np.zeros(shape, dtype=value_type)
    mem = memoryview(obj.member('buffer_'))[
        0 : int(n.prod(shape)) * np.dtype(value_type).itemsize
    ]
    c_array = np.frombuffer(mem, dtype=value_type).reshape(shape)
    # TODO: revise the memory copy of asfortranarray
    array = c_array if order == 'C' else np.asfortranarray(c_array)
    return array.view(ndarray)
//  [python]
def numpy_ndarray_resolver(obj):
    meta = obj.meta
    value_type_name = meta['value_type_']
    if value_type_name == 'object':
        view = memoryview(obj.member('buffer_'))
        return pickle.loads(view, fix_imports=True)

    value_type = normalize_dtype(value_type_name, meta.get('value_type_meta_', None))
    # process string ndarray from c++
    if value_type_name in ['str', 'string', 'std::string', 'std::__1::string']:
        from vineyard.data.arrow import binary_array_resolver
        from vineyard.data.arrow import string_array_resolver

        try:
            return string_array_resolver(obj.member('buffer_')).to_numpy(
                zero_copy_only=False
            )
        except:  # noqa: E722, pylint: disable=bare-except
            return binary_array_resolver(obj.member('buffer_')).to_numpy(
                zero_copy_only=False
            )

    shape = from_json(meta['shape_'])
    if 'order_' in meta:
        order = from_json(meta['order_'])
    else:
        order = 'C'
    if np.prod(shape) == 0:
        return np.zeros(shape, dtype=value_type)
    mem = memoryview(obj.member('buffer_'))[
        0 : int(np.prod(shape)) * np.dtype(value_type).itemsize
    ]
    c_array = np.frombuffer(mem, dtype=value_type).reshape(shape)
    # TODO: revise the memory copy of asfortranarray
    array = c_array if order == 'C' else np.asfortranarray(c_array)
    return array.view(ndarray)
//  [python]
def protective_squeeze(x: numpy.ndarray):
    """
    Squeeze dim only if it's not the last dim.
    Image dim expected to be *, H, W, C
    """
    img_shape = x.shape[-3:]
    if len(x.shape) > 3:
        n_imgs = reduce(lambda x, y: x*y, x.shape[:-3])
        if n_imgs > 1:
            img_shape = (-1,) + img_shape
    return x.reshape(img_shape)
//  [python]
def protective_squeeze(x: numpy.ndarray):
    """
    Squeeze dim only if it's not the last dim.
    Image dim expected to be *, H, W, C
    """
    img_shape = x.shape[-3:]
    if len(x.shape) > 3:
        n_imgs = numpy.prod(x.shape[:-3])
        if n_imgs > 1:
            img_shape = (-1,) + img_shape
    return x.reshape(img_shape)
//  [python]
def test_transpose_flattened_array(num_dimensions):
    np.random.seed(0)
    for _ in range(10):
        shape = np.random.randint(1, 5, (num_dimensions,)).tolist()
        axes = np.random.permutation(num_dimensions).tolist()
        volume = reduce(lambda x, y: x * y, shape)
        A = np.random.permutation(volume)
        want = np.transpose(A.reshape(shape), axes)
        got = linalg.transpose_flattened_array(A, shape, axes).reshape(want.shape)
        assert np.array_equal(want, got)
        got = linalg.transpose_flattened_array(A.reshape(shape), shape, axes).reshape(want.shape)
        assert np.array_equal(want, got)
//  [python]
def test_transpose_flattened_array(num_dimensions):
    np.random.seed(0)
    for _ in range(10):
        shape = np.random.randint(1, 5, (num_dimensions,)).tolist()
        axes = np.random.permutation(num_dimensions).tolist()
        volume = np.prod(shape)
        A = np.random.permutation(volume)
        want = np.transpose(A.reshape(shape), axes)
        got = linalg.transpose_flattened_array(A, shape, axes).reshape(want.shape)
        assert np.array_equal(want, got)
        got = linalg.transpose_flattened_array(A.reshape(shape), shape, axes).reshape(want.shape)
        assert np.array_equal(want, got)
//  [python]
def intersection(self, other: AABB) -> float:
        """
        returns the intersection volume of self and other
        """
        # pose must be identity
        assert (self.pose.position == np.array((0, 0, 0))).all() and (
            self.pose.orientation == euler.euler2quat(0, 0, 0)
        ).all()
        assert (other.pose.position == np.array((0, 0, 0))).all() and (
            other.pose.orientation == euler.euler2quat(0, 0, 0)
        ).all()
        if self.center == other.center and self.size == other.size:
            return self.volume
        min_extent = np.maximum(
            np.array(self.center) - np.array(self.size) / 2,
            np.array(other.center) - np.array(other.size) / 2,
        )
        max_extent = np.minimum(
            np.array(self.center) + np.array(self.size) / 2,
            np.array(other.center) + np.array(other.size) / 2,
        )
        intersection_size = np.maximum(max_extent - min_extent, 0)
        intersection_volume = float(np.multiply.reduce(intersection_size))
        return intersection_volume
//  [python]
def intersection(self, other: AABB) -> float:
        """
        returns the intersection volume of self and other
        """
        # pose must be identity
        assert (self.pose.position == np.array((0, 0, 0))).all() and (
            self.pose.orientation == euler.euler2quat(0, 0, 0)
        ).all()
        assert (other.pose.position == np.array((0, 0, 0))).all() and (
            other.pose.orientation == euler.euler2quat(0, 0, 0)
        ).all()
        if self.center == other.center and self.size == other.size:
            return self.volume
        min_extent = np.maximum(
            np.array(self.center) - np.array(self.size) / 2,
            np.array(other.center) - np.array(other.size) / 2,
        )
        max_extent = np.minimum(
            np.array(self.center) + np.array(self.size) / 2,
            np.array(other.center) + np.array(other.size) / 2,
        )
        intersection_size = np.maximum(max_extent - min_extent, 0)
        intersection_volume = float(np.prod(intersection_size))
        return intersection_volume
//  [python]
def zeros_aligned(shape, dtype, order='C', align=128):
    """Like `numpy.zeros()`, but the array will be aligned at `align` byte boundary."""
    nbytes = numpy.prod(shape, dtype=numpy.int64) * numpy.dtype(dtype).itemsize
    buffer = numpy.zeros(nbytes + align, dtype=numpy.uint8)  # problematic on win64 ("maximum allowed dimension exceeded")
    start_index = -buffer.ctypes.data % align
    return buffer[start_index : start_index + nbytes].view(dtype).reshape(shape, order=order)
//  [python]
def ndarange(
    shape: Sequence[int], dtype: Union[tf.DType, np.dtype] = np.float32
) -> np.ndarray:
    """np.ndarange for arbitrary input shapes."""
    # pytype doesn't understand the ternary with tf as "Any"
    dtype = (
        dtype.as_numpy_dtype if isinstance(dtype, tf.DType) else dtype
    )  # pytype: disable=attribute-error
    return np.arange(np.multiply.reduce(shape), dtype=dtype).reshape(shape)
//  [python]
def ndarange(
    shape: Sequence[int], dtype: Union[tf.DType, np.dtype] = np.float32
) -> np.ndarray:
    """np.ndarange for arbitrary input shapes."""
    # pytype doesn't understand the ternary with tf as "Any"
    dtype = (
        dtype.as_numpy_dtype if isinstance(dtype, tf.DType) else dtype
    )  # pytype: disable=attribute-error
    return np.arange(np.prod(shape), dtype=dtype).reshape(shape)
//  [python]
def mask_to_2d(mask):
    # todo todocument
    if mask.ndim > 2:
        _, d1, d2 = np.shape(mask)
        dims = d1, d2
        return scipy.sparse.coo_matrix(np.reshape(mask[:].transpose([1, 2, 0]), (
            np.multiply.reduce(dims),
            -1,
        ), order='F'))
    else:
        dims = np.shape(mask)
        return scipy.sparse.coo_matrix(np.reshape(mask, (
            np.multiply.reduce(dims),
            -1,
        ), order='F'))
//  [python]
def mask_to_2d(mask):
    # todo todocument
    if mask.ndim > 2:
        _, d1, d2 = np.shape(mask)
        dims = d1, d2
        return scipy.sparse.coo_matrix(np.reshape(mask[:].transpose([1, 2, 0]), (
            np.prod(dims),
            -1,
        ), order='F'))
    else:
        dims = np.shape(mask)
        return scipy.sparse.coo_matrix(np.reshape(mask, (
            np.prod(dims),
            -1,
        ), order='F'))
//  [python]
def memory(self) -> float:
        """Number of Bytes that the tensor takes in the RAM."""
        return np.multiply.reduce(self.shape) * self.itemsize
//  [python]
def memory(self) -> float:
        """Number of Bytes that the tensor takes in the RAM."""
        return np.prod(self.shape) * self.itemsize
//  [python]
def norm_nrg(a_):

    a = a_.copy()
    dims = a.shape
    a = a.reshape(-1, order='F')
    indx = np.argsort(a, axis=None)[::-1]
    cumEn = np.cumsum(a.flatten()[indx]**2)
    cumEn /= cumEn[-1]
    a = np.zeros(np.multiply.reduce(dims))
    a[indx] = cumEn
    return a.reshape(dims, order='F')
//  [python]
def norm_nrg(a_):

    a = a_.copy()
    dims = a.shape
    a = a.reshape(-1, order='F')
    indx = np.argsort(a, axis=None)[::-1]
    cumEn = np.cumsum(a.flatten()[indx]**2)
    cumEn /= cumEn[-1]
    a = np.zeros(np.prod(dims))
    a[indx] = cumEn
    return a.reshape(dims, order='F')
//  [python]
def get_distance_from_A(masks_gt, masks_comp, min_dist=10) -> list:
    # todo todocument

    _, d1, d2 = np.shape(masks_gt)
    dims = d1, d2
    A_ben = scipy.sparse.csc_matrix(np.reshape(masks_gt[:].transpose([1, 2, 0]), (
        n.prod(dims),
        -1,
    ), order='F'))
    A_cnmf = scipy.sparse.csc_matrix(np.reshape(masks_comp[:].transpose([1, 2, 0]), (
        n.prod(dims),
        -1,
    ), order='F'))

    cm_ben = [scipy.ndimage.center_of_mass(mm) for mm in masks_gt]
    cm_cnmf = [scipy.ndimage.center_of_mass(mm) for mm in masks_comp]

    return distance_masks([A_ben, A_cnmf], [cm_ben, cm_cnmf], min_dist)
//  [python]
def get_distance_from_A(masks_gt, masks_comp, min_dist=10) -> list:
    # todo todocument

    _, d1, d2 = np.shape(masks_gt)
    dims = d1, d2
    A_ben = scipy.sparse.csc_matrix(np.reshape(masks_gt[:].transpose([1, 2, 0]), (
        np.prod(dims),
        -1,
    ), order='F'))
    A_cnmf = scipy.sparse.csc_matrix(np.reshape(masks_comp[:].transpose([1, 2, 0]), (
        np.prod(dims),
        -1,
    ), order='F'))

    cm_ben = [scipy.ndimage.center_of_mass(mm) for mm in masks_gt]
    cm_cnmf = [scipy.ndimage.center_of_mass(mm) for mm in masks_comp]

    return distance_masks([A_ben, A_cnmf], [cm_ben, cm_cnmf], min_dist)
//  [python]
def is_data_cachable(self):
        if self.cache_enabled:
            tensor_meta = self.tensor_meta
            return (
                self.chunk_class == UncompressedChunk
                and tensor_meta.htype not in ["text", "json", "list", "polygon", "tag"]
                and tensor_meta.max_shape
                and (tensor_meta.max_shape == tensor_meta.min_shape)
                and (n.prod(tensor_meta.max_shape) < 20)
            )
        return False
//  [python]
def is_data_cachable(self):
        if self.cache_enabled:
            tensor_meta = self.tensor_meta
            return (
                self.chunk_class == UncompressedChunk
                and tensor_meta.htype not in ["text", "json", "list", "polygon", "tag"]
                and tensor_meta.max_shape
                and (tensor_meta.max_shape == tensor_meta.min_shape)
                and (np.prod(tensor_meta.max_shape) < 20)
            )
        return False
//  [python]
def detect_duplicates(file_name:str, dist_thr:float = 0.1, FOV:tuple[int, ...] = (512, 512)) -> tuple[list, list]:
    """
    Removes duplicate ROIs from file file_name

    Args:
        file_name:  .zip file with all rois

        dist_thr:   distance threshold for duplicate detection

        FOV:        dimensions of the FOV

    Returns:
        duplicates  : list of indices with duplicate entries

        ind_keep    : list of kept indices

    """
    rois = nf_read_roi_zip(file_name, FOV)
    cm = [scipy.ndimage.center_of_mass(mm) for mm in rois]
    sp_rois = scipy.sparse.csc_matrix(np.reshape(rois, (rois.shape[0], np.prod(FOV))).T)
    D = distance_masks([sp_rois, sp_rois], [cm, cm], 10)[0]
    np.fill_diagonal(D, 1)
    indices = np.where(D < dist_thr)   # pairs of duplicate indices

    ind = list(np.unique(indices[1][indices[1] > indices[0]]))
    ind_keep = list(set(range(D.shape[0])) - set(ind))
    duplicates = list(np.unique(np.concatenate((indices[0], indices[1]))))

    return duplicates, ind_keep
//  [python]
def tune():

    with open('stencil.cu', 'r') as f:
        kernel_string = f.read()

    problem_size = (4096, 2048)
    size = numpy.prod(problem_size)

    x_old = numpy.random.randn(size).astype(numpy.float32)
    x_new = numpy.copy(x_old)
    args = [x_new, x_old]

    tune_params = OrderedDict()
    tune_params["block_size_x"] = [32*i for i in range(1,9)]
    tune_params["block_size_y"] = [2**i for i in range(6)]

    grid_div_x = ["block_size_x"]
    grid_div_y = ["block_size_y"]

    return kernel_tuner.tune_kernel("stencil_kernel", kernel_string, problem_size,
        args, tune_params, grid_div_x=grid_div_x, grid_div_y=grid_div_y,
        verbose = True)
//  [python]
def _test(self, input_shape, dtype, **kwargs):
    # Use negative numbers to make sure there isn't any zero padding getting
    # used.
    x = -np.arange(np.prod(input_shape), dtype=dtype).reshape(input_shape) - 1
    y1 = pool_direct(input=x, **kwargs)
    y2 = nn_ops.pool(input=x, **kwargs)
    self.assertAllClose(y1, self.evaluate(y2), rtol=1e-2, atol=1e-2)
//  [python]
def _test_gradient(self, input_shape, dtype, **kwargs):
    x_val = (
        -np.arange(np.prod(input_shape), dtype=dtype).reshape(input_shape) - 1
    )
    x = constant_op.constant(x_val, name="x", dtype=dtype)
    output = nn_ops.pool(input=x, **kwargs)
    y_shape = output.get_shape().as_list()
    err = gradient_checker.compute_gradient_error(
        [x], [input_shape], output, y_shape, x_init_value=[x_val]
    )
    err_tolerance = 1e-2
    if dtype == dtypes.float16:
      err_tolerance = 1.1
      # TODO: this is too high.
      # investigate precision issues.
    self.assertLess(err, err_tolerance)
//  [python]
def _makeIncremental(self, shape, dtype):
    data = np.arange(np.multiply.reduce(shape)).reshape(shape).astype(dtype.as_numpy_dtype)
    if dtype.is_complex:
      data -= 2j * data
    return data
//  [python]
def _makeIncremental(self, shape, dtype):
    data = np.arange(np.prod(shape)).reshape(shape).astype(dtype.as_numpy_dtype)
    if dtype.is_complex:
      data -= 2j * data
    return data
//  [python]
def testBCast(self):
    shapes = [
        ([1, 3, 2], [1]),
        ([1, 3, 2], [2]),
        ([1, 3, 2], [3, 2]),
        ([1, 3, 2], [3, 1]),
        ([1, 3, 2], [1, 3, 2]),
        ([1, 3, 2], [2, 3, 1]),
        ([1, 3, 2], [2, 1, 1]),
        ([1, 3, 2], [1, 3, 1]),
        ([2, 1, 5], [2, 3, 1]),
        ([2, 0, 5], [2, 0, 1]),
        ([2, 3, 0], [2, 3, 1]),
    ]
    for xs, ys in shapes:
      x = np.random.randint(0, 2, np.prod(xs)).astype(np.bool).reshape(xs)
      y = np.random.randint(0, 2, np.prod(ys)).astype(np.bool).reshape(ys)
      for use_gpu in [True, False]:
        self._compareBinary(x, y, np.logical_and, math_ops.logical_and, use_gpu)
        self._compareBinary(x, y, np.logical_or, math_ops.logical_or, use_gpu)
        self._compareBinary(x, y, np.logical_xor, math_ops.logical_xor, use_gpu)
//  [python]
def count_model_params(self, net: nn.HybridBlock) -> int:
        params = net.collect_params()
        num_params = 0
        for p in params:
            v = params[p]
            num_params += n.prod(v.shape)
        return num_params
//  [python]
def count_model_params(self, net: nn.HybridBlock) -> int:
        params = net.collect_params()
        num_params = 0
        for p in params:
            v = params[p]
            num_params += np.prod(v.shape)
        return num_params
//  [python]
def size(self) -> int:
        """Returns the number of elements in the tensor dataclass batch dimension."""
        if len(self._shape) == 0:
            return 1
        return int(np.prod(self._shape))
//  [python]
def hypeIndicatorExact(points, bounds, k):
    """HypE algorithm. Python implementation of the Matlab code available at
    https://sop.tik.ee.ethz.ch/download/supplementary/hype/

    Args:
        points(array): 2D array containing the objective values of the
        population
        bounds(array): 1D array containing the reference point from which to
        compute the hyper-volume
        k(int): HypE parameter
    """

    Ps = points.shape[0]
    if k < 0:
        k = Ps
    actDim = points.shape[1]
    pvec = numpy.arange(points.shape[0])

    alpha = []
    for i in range(1, k + 1):
        j = numpy.arange(1, i)
        alpha.append(n.prod((k - j) / (Ps - j) / i))
    alpha = numpy.asarray(alpha)

    return hypesub(points.shape[0], points, actDim, bounds, pvec, alpha, k)
//  [python]
def hypeIndicatorExact(points, bounds, k):
    """HypE algorithm. Python implementation of the Matlab code available at
    https://sop.tik.ee.ethz.ch/download/supplementary/hype/

    Args:
        points(array): 2D array containing the objective values of the
        population
        bounds(array): 1D array containing the reference point from which to
        compute the hyper-volume
        k(int): HypE parameter
    """

    Ps = points.shape[0]
    if k < 0:
        k = Ps
    actDim = points.shape[1]
    pvec = numpy.arange(points.shape[0])

    alpha = []
    for i in range(1, k + 1):
        j = numpy.arange(1, i)
        alpha.append(numpy.prod((k - j) / (Ps - j) / i))
    alpha = numpy.asarray(alpha)

    return hypesub(points.shape[0], points, actDim, bounds, pvec, alpha, k)
//  [python]
def count_parameters(module: torch.nn.Module) -> int:
    return int(sum(n.prod(p.shape) for p in module.parameters()))
//  [python]
def count_parameters(module: torch.nn.Module) -> int:
    return int(sum(np.prod(p.shape) for p in module.parameters()))
//  [python]
def hypeIndicatorSampled(points, bounds, k, nrOfSamples):
    """Monte-Carlo approximation of the HypE algorithm. Python implementation
    of the Matlab code available at
    https://sop.tik.ee.ethz.ch/download/supplementary/hype/

    Args:
        points(array): 2D array containing the objective values of the
        population
        bounds(array): 1D array containing the reference point from which to
        compute the hyper-volume
        k(int): HypE parameter
        nrOfSamples(int): number of random samples to use for the
        Monte-Carlo approximation
    """

    nrP = points.shape[0]
    dim = points.shape[1]
    F = numpy.zeros(nrP)

    BoxL = numpy.min(points, axis=0)

    alpha = []
    for i in range(1, k + 1):
        j = numpy.arange(1, i)
        alpha.append(numpy.prod((k - j) / (nrP - j) / i))
    alpha = numpy.asarray(alpha + [0.0] * nrP)

    S = numpy.random.uniform(low=BoxL, high=bounds, size=(nrOfSamples, dim))

    dominated = numpy.zeros(nrOfSamples, dtype="uint")
    for j in range(1, nrP + 1):
        B = S - points[j - 1]
        ind = numpy.sum(B >= 0, axis=1) == dim
        dominated[ind] += 1

    for j in range(1, nrP + 1):
        B = S - points[j - 1]
        ind = numpy.sum(B >= 0, axis=1) == dim
        x = dominated[ind]
        F[j - 1] = numpy.sum(alpha[x - 1])

    F = F * n.prod(bounds - BoxL) / nrOfSamples

    return F
//  [python]
def hypeIndicatorSampled(points, bounds, k, nrOfSamples):
    """Monte-Carlo approximation of the HypE algorithm. Python implementation
    of the Matlab code available at
    https://sop.tik.ee.ethz.ch/download/supplementary/hype/

    Args:
        points(array): 2D array containing the objective values of the
        population
        bounds(array): 1D array containing the reference point from which to
        compute the hyper-volume
        k(int): HypE parameter
        nrOfSamples(int): number of random samples to use for the
        Monte-Carlo approximation
    """

    nrP = points.shape[0]
    dim = points.shape[1]
    F = numpy.zeros(nrP)

    BoxL = numpy.min(points, axis=0)

    alpha = []
    for i in range(1, k + 1):
        j = numpy.arange(1, i)
        alpha.append(numpy.prod((k - j) / (nrP - j) / i))
    alpha = numpy.asarray(alpha + [0.0] * nrP)

    S = numpy.random.uniform(low=BoxL, high=bounds, size=(nrOfSamples, dim))

    dominated = numpy.zeros(nrOfSamples, dtype="uint")
    for j in range(1, nrP + 1):
        B = S - points[j - 1]
        ind = numpy.sum(B >= 0, axis=1) == dim
        dominated[ind] += 1

    for j in range(1, nrP + 1):
        B = S - points[j - 1]
        ind = numpy.sum(B >= 0, axis=1) == dim
        x = dominated[ind]
        F[j - 1] = numpy.sum(alpha[x - 1])

    F = F * numpy.prod(bounds - BoxL) / nrOfSamples

    return F
//  [python]
def shape_normalize(image, batch_axes=None, channel_axes=None):
        """Permutates batch dimensions to the front and collapse into
        one dimension. Resulting array has to be in the form of (B,...)
        where B is the number of batch dimensions.

        Parameters
        ----------
        image : numpy.ArrayLike
        batch_axes : list
        channel_axes : list

        Returns
        -------
        Tuple of normalized_image, axes_permutation, permutated_image.shape : tuple

        """
        if batch_axes is None:
            batch_axes = (False,) * len(image.shape)
        if channel_axes is None:
            channel_axes = (False,) * len(image.shape)

        # Singleton dimensions are automatically batch dimension, unless it is a channel dimension, trivially.
        batch_axes = tuple(
            True if s == 1 and not c else b
            for b, c, s in zip(batch_axes, channel_axes, image.shape)
        )

        # get indices for different types of dimensions and their length
        batch_indices = [index for index, value in enumerate(batch_axes) if value]
        batch_length = int(reduce(lambda x, y: x * y, [image.shape[index] for index in batch_indices]))

        channel_indices = (
            [index for index, value in enumerate(channel_axes) if value]
            if channel_axes
            else []
        )
        channel_length = int(
            reduce(lambda x, y: x * y, [image.shape[index] for index in channel_indices])
        )

        spacetime_indices = [
            index
            for index in range(len(image.shape))
            if index not in batch_indices + channel_indices
        ]

        # Axes permutation
        axes_permutation = batch_indices + channel_indices + spacetime_indices

        # Bring all batch dimensions to front
        permutated_image = numpy.transpose(image, axes_permutation)

        # Collapse batch dimensions into one, same for channel dimensions
        spacetime_shape = tuple([image.shape[i] for i in spacetime_indices])
        normalized_shape = (batch_length, channel_length) + spacetime_shape
        normalized_shape = tuple(
            s for i, s in enumerate(normalized_shape) if s != 1 or i <= 1
        )

        # Reshape array:
        normalized_image = permutated_image.reshape(normalized_shape)

        return (normalized_image, axes_permutation, permutated_image.shape)
//  [python]
def shape_normalize(image, batch_axes=None, channel_axes=None):
        """Permutates batch dimensions to the front and collapse into
        one dimension. Resulting array has to be in the form of (B,...)
        where B is the number of batch dimensions.

        Parameters
        ----------
        image : numpy.ArrayLike
        batch_axes : list
        channel_axes : list

        Returns
        -------
        Tuple of normalized_image, axes_permutation, permutated_image.shape : tuple

        """
        if batch_axes is None:
            batch_axes = (False,) * len(image.shape)
        if channel_axes is None:
            channel_axes = (False,) * len(image.shape)

        # Singleton dimensions are automatically batch dimension, unless it is a channel dimension, trivially.
        batch_axes = tuple(
            True if s == 1 and not c else b
            for b, c, s in zip(batch_axes, channel_axes, image.shape)
        )

        # get indices for different types of dimensions and their length
        batch_indices = [index for index, value in enumerate(batch_axes) if value]
        batch_length = int(numpy.prod([image.shape[index] for index in batch_indices]))

        channel_indices = (
            [index for index, value in enumerate(channel_axes) if value]
            if channel_axes
            else []
        )
        channel_length = int(
            numpy.prod([image.shape[index] for index in channel_indices])
        )

        spacetime_indices = [
            index
            for index in range(len(image.shape))
            if index not in batch_indices + channel_indices
        ]

        # Axes permutation
        axes_permutation = batch_indices + channel_indices + spacetime_indices

        # Bring all batch dimensions to front
        permutated_image = numpy.transpose(image, axes_permutation)

        # Collapse batch dimensions into one, same for channel dimensions
        spacetime_shape = tuple([image.shape[i] for i in spacetime_indices])
        normalized_shape = (batch_length, channel_length) + spacetime_shape
        normalized_shape = tuple(
            s for i, s in enumerate(normalized_shape) if s != 1 or i <= 1
        )

        # Reshape array:
        normalized_image = permutated_image.reshape(normalized_shape)

        return (normalized_image, axes_permutation, permutated_image.shape)
//  [python]
def _get_indices_dict(label_list, keys):
    shape = list(map(len, keys))

    group_index = get_group_index(label_list, shape, sort=True, xnull=True)
    ngroups = ((group_index.size and group_index.max()) + 1) \
        if _int64_overflow_possible(shape) \
        else n.prod(shape, dtype='i8')

    sorter = _get_group_index_sorter(group_index, ngroups)

    sorted_labels = [lab.take(sorter) for lab in label_list]
    group_index = group_index.take(sorter)

    return lib.indices_fast(sorter, group_index, keys, sorted_labels)