jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def remove_infs_from_variable(v):
    # 32 - bit floating point: torch.FloatTensor, torch.cuda.FloatTensor
    # 64 - bit floating point: torch.DoubleTensor, torch.cuda.DoubleTensor
    # 16 - bit floating point: torch.HalfTensor, torch.cuda.HalfTensor

    # todo: maybe find a cleaner way of handling this
    # this is to make sure that subsequent sums work (hence will be smaller than it could be,
    # but values of this size should not occur in practice anyway
    sz = v.size()
    reduction_factor = np.prod(np.array(sz))
    condition = True

    if type(v.data) == torch.cuda.FloatTensor or v.data.dtype==torch.float32:
        return torch.clamp(v,
                           min=(np.asscalar(np.finfo('float32').min))/reduction_factor,
                           max=(np.asscalar(np.finfo('float32').max))/reduction_factor)
    elif v.data.dtype == torch.DoubleTensor or type(v.data) == torch.cuda.DoubleTensor:
        return torch.clamp(v,
                           min=(np.asscalar(np.finfo('float64').min))/reduction_factor,
                           max=(np.asscalar(np.finfo('float64').max))/reduction_factor)
    elif v.data.dtype == torch.HalfTensor or type(v.data) == torch.cuda.HalfTensor:
        return torch.clamp(v,
                           min=(np.asscalar(np.finfo('float16').min))/reduction_factor,
                           max=(np.asscalar(np.finfo('float16').max))/reduction_factor)
    else:
        raise ValueError('Unknown data type: ' + str( type(v.data)))
//  [python]
def test_concatenations_and_stacking():
    for backend in imp_op_backends:
        print("testing shapes for ", backend.framework_name)
        for n_arrays in [1, 2, 5]:
            shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
            for shape in shapes:
                arrays1 = [numpy.arange(i, i + numpy.prod(shape)).reshape(shape) for i in range(n_arrays)]
                arrays2 = [backend.from_numpy(array) for array in arrays1]
                result0 = numpy.asarray(arrays1)
                result1 = rearrange(arrays1, "...->...")
                result2 = rearrange(arrays2, "...->...")
                assert numpy.array_equal(result0, result1)
                assert numpy.array_equal(result1, backend.to_numpy(result2))

                result1 = rearrange(arrays1, "b ... -> ... b")
                result2 = rearrange(arrays2, "b ... -> ... b")
                assert numpy.array_equal(result1, backend.to_numpy(result2))
//  [python]
def test_reduction_stress_imperatives():
    for backend in imp_op_backends:
        print("Stress-testing reduction for ", backend.framework_name)
        for reduction in REDUCTIONS + ("rearrange",):
            dtype = "int64"
            coincide = numpy.array_equal
            if reduction in ["mean", "prod"]:
                dtype = "float64"
                coincide = numpy.allclose
            max_dim = 11
            if "oneflow" in backend.framework_name:
                max_dim = 7
            if "paddle" in backend.framework_name:
                max_dim = 9
            for n_axes in range(max_dim):
                shape = numpy.random.randint(2, 4, size=n_axes)
                permutation = numpy.random.permutation(n_axes)
                skipped = 0 if reduction == "rearrange" else numpy.random.randint(n_axes + 1)
                left = " ".join("x" + str(i) for i in range(n_axes))
                right = " ".join("x" + str(i) for i in permutation[skipped:])
                pattern = left + "->" + right
                x = numpy.arange(1, 1 + numpy.prod(shape), dtype=dtype).reshape(shape)
                if reduction == "prod":
                    x /= x.mean()  # to avoid overflows
                result1 = reduce(x, pattern, reduction=reduction)
                result2 = x.transpose(permutation)
                if skipped > 0:
                    result2 = getattr(result2, reduction)(axis=tuple(range(skipped)))
                assert coincide(result1, result2)
                check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False)
//  [python]
def interpolate_indices(self, *params, threshold=1.e-6):
        """Interpolate the underlying SPECTRA extension at a given set of input
        parameters.

        This performs a simple linear interpolation, where the first parameter
        (the actual model) is treated in a special fashion in that is not
        interpolated at all---in fact the function requires that the value
        passed as an argument to the function is one of the four valis values.

        For all the other parameters, the corresponding grids of tabulated values
        are bisected at the target value and the two nearest elements are
        weighted proportinally to the relative distance of the target.
        """
        model, *params = params
        model = int(model)
        assert xMagnetarModelsT2020.has_value(model)
        par_values = self.par_values[1:-1]
        indices = []
        weights = []
        # Bisect the parameter grids.
        for value, grid in zip(params, par_values):
            i = numpy.searchsorted(grid, value)
            indices.append((i - 1, i))
            lo, hi = grid[i - 1], grid[i]
            w = (hi - value) / (hi - lo)
            # Small hack to avoid numerical artifacts due to the fact that the
            # parameter grids are saved in single precision, while everywhere
            # else we are operating in double precision. Essentially if the
            # weights are too close to either 0 or 1, we attribute that to simple
            # numerical rounding.
            if w < threshold:
                w = 0.
            elif 1. - w < threshold:
                w = 1.
            weights.append((w, 1. - w))
        # Create all the combination of the relevant parameters and weights.
        indices = self._combinations(indices)
        weights = self._combinations(weights)
        # Multiply the weights along the proper axis to come up with one value
        # for each parameter set.
        weights = numpy.multiply.reduce(weights, axis=1)
        # Get rid of the combinations that will not contribute to the weighted
        # average.
        mask = weights > 0.
        indices = indices[mask]
        weights = weights[mask]
        # Prepend the model to the indices.
        model_indices = numpy.full((len(indices), 1), model - 1)
        indices = numpy.hstack((model_indices, indices))
        return indices, weights
//  [python]
def interpolate_indices(self, *params, threshold=1.e-6):
        """Interpolate the underlying SPECTRA extension at a given set of input
        parameters.

        This performs a simple linear interpolation, where the first parameter
        (the actual model) is treated in a special fashion in that is not
        interpolated at all---in fact the function requires that the value
        passed as an argument to the function is one of the four valis values.

        For all the other parameters, the corresponding grids of tabulated values
        are bisected at the target value and the two nearest elements are
        weighted proportinally to the relative distance of the target.
        """
        model, *params = params
        model = int(model)
        assert xMagnetarModelsT2020.has_value(model)
        par_values = self.par_values[1:-1]
        indices = []
        weights = []
        # Bisect the parameter grids.
        for value, grid in zip(params, par_values):
            i = numpy.searchsorted(grid, value)
            indices.append((i - 1, i))
            lo, hi = grid[i - 1], grid[i]
            w = (hi - value) / (hi - lo)
            # Small hack to avoid numerical artifacts due to the fact that the
            # parameter grids are saved in single precision, while everywhere
            # else we are operating in double precision. Essentially if the
            # weights are too close to either 0 or 1, we attribute that to simple
            # numerical rounding.
            if w < threshold:
                w = 0.
            elif 1. - w < threshold:
                w = 1.
            weights.append((w, 1. - w))
        # Create all the combination of the relevant parameters and weights.
        indices = self._combinations(indices)
        weights = self._combinations(weights)
        # Multiply the weights along the proper axis to come up with one value
        # for each parameter set.
        weights = numpy.prod(weights, 1)
        # Get rid of the combinations that will not contribute to the weighted
        # average.
        mask = weights > 0.
        indices = indices[mask]
        weights = weights[mask]
        # Prepend the model to the indices.
        model_indices = numpy.full((len(indices), 1), model - 1)
        indices = numpy.hstack((model_indices, indices))
        return indices, weights
//  [python]
def from_arrow(array: pa.Array | pa.ChunkedArray, name: str = "arrow_series") -> Series:
        """
        Construct a Series from an pyarrow array or chunked array.

        Args:
            array: The pyarrow (chunked) array whose data we wish to put in the Series.
            name: The name associated with the Series; this is usually the column name.
        """
        if DataType.from_arrow_type(array.type) == DataType.python():
            # If the Arrow type is not natively supported, go through the Python list path.
            return Series.from_pylist(array.to_pylist(), name=name, pyobj="force")
        elif isinstance(array, pa.Array):
            array = ensure_array(array)
            if _RAY_DATA_EXTENSIONS_AVAILABLE and isinstance(array.type, ArrowTensorType):
                storage_series = Series.from_arrow(array.storage, name=name)
                series = storage_series.cast(
                    DataType.fixed_size_list(
                        DataType.from_arrow_type(array.type.scalar_type), int(np.prod(array.type.shape))
                    )
                )
                return series.cast(DataType.from_arrow_type(array.type))
            elif _RAY_DATA_EXTENSIONS_AVAILABLE and isinstance(array.type, ArrowVariableShapedTensorType):
                return Series.from_numpy(array.to_numpy(zero_copy_only=False), name=name)
            elif isinstance(array.type, getattr(pa, "FixedShapeTensorType", ())):
                series = Series.from_arrow(array.storage, name=name)
                return series.cast(DataType.from_arrow_type(array.type))
            else:
                pys = PySeries.from_arrow(name, array)
                return Series._from_pyseries(pys)
        elif isinstance(array, pa.ChunkedArray):
            array = ensure_chunked_array(array)
            arr_type = array.type
            if isinstance(arr_type, pa.BaseExtensionType):
                combined_storage_array = array.cast(arr_type.storage_type).combine_chunks()
                combined_array = arr_type.wrap_array(combined_storage_array)
            else:
                combined_array = array.combine_chunks()
            return Series.from_arrow(combined_array)
        else:
            raise TypeError(f"expected either PyArrow Array or Chunked Array, got {type(array)}")
//  [python]
def test_rearrange_consistency_numpy():
    shape = [1, 2, 3, 5, 7, 11]
    x = numpy.arange(numpy.prod(shape)).reshape(shape)
    for pattern in [
        "a b c d e f -> a b c d e f",
        "b a c d e f -> a b d e f c",
        "a b c d e f -> f e d c b a",
        "a b c d e f -> (f e) d (c b a)",
        "a b c d e f -> (f e d c b a)",
    ]:
        result = rearrange(x, pattern)
        assert len(numpy.setdiff1d(x, result)) == 0
        assert result.dtype == x.dtype

    result = rearrange(x, "a b c d e f -> a (b) (c d e) f")
    assert numpy.array_equal(x.flatten(), result.flatten())

    result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11")
    assert numpy.array_equal(x, result)

    result1 = rearrange(x, "a b c d e f -> f e d c b a")
    result2 = rearrange(x, "f e d c b a -> a b c d e f")
    assert numpy.array_equal(result1, result2)

    result = rearrange(rearrange(x, "a b c d e f -> (f d) c (e b) a"), "(f d) c (e b) a -> a b c d e f", b=2, d=5)
    assert numpy.array_equal(x, result)

    sizes = dict(zip("abcdef", shape))
    temp = rearrange(x, "a b c d e f -> (f d) c (e b) a", **sizes)
    result = rearrange(temp, "(f d) c (e b) a -> a b c d e f", **sizes)
    assert numpy.array_equal(x, result)

    x2 = numpy.arange(2 * 3 * 4).reshape([2, 3, 4])
    result = rearrange(x2, "a b c -> b c a")
    assert x2[1, 2, 3] == result[2, 3, 1]
    assert x2[0, 1, 2] == result[1, 2, 0]
//  [python]
def test_array(self, arr):
        assert allclose(n.prod(arr), num.prod(arr))
//  [python]
def test_array(self, arr):
        assert allclose(np.prod(arr), num.prod(arr))
//  [python]
def setUp(self):
        self.x = numpy.arange(numpy.multiply.reduce(self.in_shape), dtype=numpy.float32)
        self.x = self.x.reshape(self.in_shape)

        out_row = numpy.linspace(0, 123 - 1, 345, dtype=numpy.float32)
        out_col = numpy.linspace(0, (987 - 1) * 123, 765, dtype=numpy.float32)
        self.out = (out_row + out_col[:, None]).reshape(self.output_shape)
//  [python]
def setUp(self):
        self.x = numpy.arange(numpy.prod(self.in_shape), dtype=numpy.float32)
        self.x = self.x.reshape(self.in_shape)

        out_row = numpy.linspace(0, 123 - 1, 345, dtype=numpy.float32)
        out_col = numpy.linspace(0, (987 - 1) * 123, 765, dtype=numpy.float32)
        self.out = (out_row + out_col[:, None]).reshape(self.output_shape)
//  [python]
def generate_input_data(
    ids_shape,
    vocab_size,
    hidden_size,
    weight_dtype,
    ids_dtype,
    allow_duplicate_id=True,
    rank=None,
    nranks=None,
    allow_pure_random=False,
):
    max_id = vocab_size if rank is None else vocab_size * nranks
    if allow_duplicate_id:
        ids = np.random.randint(low=0, high=max_id, size=ids_shape)
    else:
        sequence = list(range(max_id))
        numel = int(np.prod(ids_shape))
        if len(sequence) < numel:
            return None, None, None
        ids = np.array(random.sample(sequence, numel)).reshape(ids_shape)

    ids = paddle.to_tensor(ids).astype(ids_dtype)
    ids.stop_gradient = True

    weight = paddle.randn([vocab_size, hidden_size]).astype(weight_dtype)
    weight.stop_gradient = False

    out_grad_shape = list(ids_shape) + [hidden_size]
    if allow_duplicate_id and not allow_pure_random:
        out_grad = paddle.randint(low=-10, high=10, shape=out_grad_shape)
    else:
        out_grad = paddle.randn(out_grad_shape)
    out_grad = out_grad.astype(weight.dtype)
    return ids, weight, out_grad
//  [python]
def test_axis_out_bound(self):
        expected_exc = np.AxisError
        arr = [-1, 0, 1, 2, 10]
        with pytest.raises(expected_exc):
            n.prod(arr, axis=2)
        with pytest.raises(expected_exc):
            num.prod(arr, axis=2)
//  [python]
def test_axis_out_bound(self):
        expected_exc = np.AxisError
        arr = [-1, 0, 1, 2, 10]
        with pytest.raises(expected_exc):
            np.prod(arr, axis=2)
        with pytest.raises(expected_exc):
            num.prod(arr, axis=2)
//  [python]
def _receive_array(kernel, embedding_map):
    num_dims = kernel._read_int8()
    shape = tuple(kernel._read_int32() for _ in range(num_dims))
    tag = chr(kernel._read_int8())
    fn = receivers[tag]
    length = n.prod(shape)
    if tag == "b":
        buffer = kernel._read(length)
        elems = numpy.ndarray((length, ), '?', buffer)
    elif tag == "i":
        buffer = kernel._read(4 * length)
        elems = numpy.ndarray((length, ), kernel.endian + 'i4', buffer)
    elif tag == "I":
        buffer = kernel._read(8 * length)
        elems = numpy.ndarray((length, ), kernel.endian + 'i8', buffer)
    elif tag == "f":
        buffer = kernel._read(8 * length)
        elems = numpy.ndarray((length, ), kernel.endian + 'd', buffer)
    else:
        fn = receivers[tag]
        elems = []
        for _ in range(n.prod(shape)):
            # discard the tag
            kernel._read_int8()
            item = fn(kernel, embedding_map)
            elems.append(item)
        elems = numpy.array(elems)
    return elems.reshape(shape)
//  [python]
def _receive_array(kernel, embedding_map):
    num_dims = kernel._read_int8()
    shape = tuple(kernel._read_int32() for _ in range(num_dims))
    tag = chr(kernel._read_int8())
    fn = receivers[tag]
    length = numpy.prod(shape)
    if tag == "b":
        buffer = kernel._read(length)
        elems = numpy.ndarray((length, ), '?', buffer)
    elif tag == "i":
        buffer = kernel._read(4 * length)
        elems = numpy.ndarray((length, ), kernel.endian + 'i4', buffer)
    elif tag == "I":
        buffer = kernel._read(8 * length)
        elems = numpy.ndarray((length, ), kernel.endian + 'i8', buffer)
    elif tag == "f":
        buffer = kernel._read(8 * length)
        elems = numpy.ndarray((length, ), kernel.endian + 'd', buffer)
    else:
        fn = receivers[tag]
        elems = []
        for _ in range(numpy.prod(shape)):
            # discard the tag
            kernel._read_int8()
            item = fn(kernel, embedding_map)
            elems.append(item)
        elems = numpy.array(elems)
    return elems.reshape(shape)
//  [python]
def matrix_to_block_tree(matrix, outer_tree, inner_tree):
    """Convert a matrix (2-dimensional array) to block-tree.

    A block tree most often arises when one applies an operation to a function that maps
    between two trees. For certain functions this results in a 2-dimensional data array.
    Two main examples are the Jacobian of the function f : inner_tree -> outer_tree,
    which results in a block tree structure, or the covariance matrix of a tree, in
    which case outer_tree = inner_tree.

    Args:
        matrix (numpy.ndarray): 2d representation of the block tree. Has shape (m, n).
        outer_tree: A pytree. If flattened to scalars has length m.
        inner_tree: A pytree. If flattened to scalars has length n.

    Returns:
        block_tree: A (block) pytree.

    """
    _check_dimensions_matrix(matrix, outer_tree, inner_tree)

    flat_outer, treedef_outer = tree_flatten(outer_tree)
    flat_inner, treedef_inner = tree_flatten(inner_tree)

    flat_outer_np = [_convert_to_numpy(leaf, only_pandas=True) for leaf in flat_outer]
    flat_inner_np = [_convert_to_numpy(leaf, only_pandas=True) for leaf in flat_inner]

    shapes_outer = [np.shape(a) for a in flat_outer_np]
    shapes_inner = [np.shape(a) for a in flat_inner_np]

    block_bounds_outer = np.cumsum([int(n.prod(s)) for s in shapes_outer[:-1]])
    block_bounds_inner = np.cumsum([int(n.prod(s)) for s in shapes_inner[:-1]])

    blocks = []
    for leaf_outer, s1, submat in zip(
        flat_outer, shapes_outer, np.split(matrix, block_bounds_outer, axis=0)
    ):
        row = []
        for leaf_inner, s2, block_values in zip(
            flat_inner, shapes_inner, np.split(submat, block_bounds_inner, axis=1)
        ):
            raw_block = block_values.reshape((*s1, *s2))
            block = _convert_raw_block_to_pandas(raw_block, leaf_outer, leaf_inner)
            row.append(block)

        blocks.append(row)

    block_tree = tree_unflatten(
        treedef_outer, [tree_unflatten(treedef_inner, row) for row in blocks]
    )

    return block_tree
//  [python]
def matrix_to_block_tree(matrix, outer_tree, inner_tree):
    """Convert a matrix (2-dimensional array) to block-tree.

    A block tree most often arises when one applies an operation to a function that maps
    between two trees. For certain functions this results in a 2-dimensional data array.
    Two main examples are the Jacobian of the function f : inner_tree -> outer_tree,
    which results in a block tree structure, or the covariance matrix of a tree, in
    which case outer_tree = inner_tree.

    Args:
        matrix (numpy.ndarray): 2d representation of the block tree. Has shape (m, n).
        outer_tree: A pytree. If flattened to scalars has length m.
        inner_tree: A pytree. If flattened to scalars has length n.

    Returns:
        block_tree: A (block) pytree.

    """
    _check_dimensions_matrix(matrix, outer_tree, inner_tree)

    flat_outer, treedef_outer = tree_flatten(outer_tree)
    flat_inner, treedef_inner = tree_flatten(inner_tree)

    flat_outer_np = [_convert_to_numpy(leaf, only_pandas=True) for leaf in flat_outer]
    flat_inner_np = [_convert_to_numpy(leaf, only_pandas=True) for leaf in flat_inner]

    shapes_outer = [np.shape(a) for a in flat_outer_np]
    shapes_inner = [np.shape(a) for a in flat_inner_np]

    block_bounds_outer = np.cumsum([int(np.prod(s)) for s in shapes_outer[:-1]])
    block_bounds_inner = np.cumsum([int(np.prod(s)) for s in shapes_inner[:-1]])

    blocks = []
    for leaf_outer, s1, submat in zip(
        flat_outer, shapes_outer, np.split(matrix, block_bounds_outer, axis=0)
    ):
        row = []
        for leaf_inner, s2, block_values in zip(
            flat_inner, shapes_inner, np.split(submat, block_bounds_inner, axis=1)
        ):
            raw_block = block_values.reshape((*s1, *s2))
            block = _convert_raw_block_to_pandas(raw_block, leaf_outer, leaf_inner)
            row.append(block)

        blocks.append(row)

    block_tree = tree_unflatten(
        treedef_outer, [tree_unflatten(treedef_inner, row) for row in blocks]
    )

    return block_tree
//  [python]
def test_keepdims(self):
        in_shape = (2, 3, 4)
        arr_num = num.random.random(in_shape) * 10
        arr_np = np.array(arr_num)
        out_np = np.prod(arr_np, axis=2, keepdims=True)
        out_num = num.prod(arr_num, axis=2, keepdims=True)
        assert allclose(out_np, out_num)
//  [python]
def test_out_negative(self):
        expected_exc = ValueError
        in_shape = (2, 3, 4)
        out_shape = (2, 3, 3)
        arr_np = np.ndarray(in_shape)
        out_np = np.ndarray(out_shape)
        arr_num = num.ndarray(in_shape)
        out_num = num.ndarray(out_shape)
        with pytest.raises(expected_exc):
            np.multiply.reduce(arr_np, axis=2, out=out_np)
        with pytest.raises(expected_exc):
            num.multiply.reduce(arr_num, axis=2, out=out_num)
//  [python]
def test_out_negative(self):
        expected_exc = ValueError
        in_shape = (2, 3, 4)
        out_shape = (2, 3, 3)
        arr_np = np.ndarray(in_shape)
        out_np = np.ndarray(out_shape)
        arr_num = num.ndarray(in_shape)
        out_num = num.ndarray(out_shape)
        with pytest.raises(expected_exc):
            np.prod(arr_np, out=out_np, axis=2)
        with pytest.raises(expected_exc):
            num.prod(arr_num, out=out_num, axis=2)
//  [python]
def test_small_images():
    """Test mixture of gaussians thresholding with few pixels

    Run MOG to see if it blows up, given 0-10 pixels"""
    r = numpy.random.RandomState()
    r.seed(91)
    image = r.uniform(size=(9, 11))
    ii, jj = numpy.mgrid[0 : image.shape[0], 0 : image.shape[1]]
    ii, jj = ii.flatten(), jj.flatten()

    for threshold_method in (
        cellprofiler.modules.threshold.TM_LI,
        cellprofiler.modules.threshold.TM_OTSU,
        cellprofiler.modules.threshold.TM_ROBUST_BACKGROUND,
    ):
        for i in range(11):
            mask = numpy.zeros(image.shape, bool)
            if i:
                p = r.permutation(numpy.multiply.reduce(image.shape))[:i]
                mask[ii[p], jj[p]] = True
            workspace, x = make_workspace(image, mask)
            x.global_operation.value = threshold_method
            x.threshold_scope.value = TS_GLOBAL
            l, g, _, _, _ = x.get_threshold(
                cellprofiler_core.image.Image(image, mask=mask), workspace
            )
            v = image[mask]
            image = r.uniform(size=(9, 11))
            image[mask] = v
            l1, g1, _, _, _ = x.get_threshold(
                cellprofiler_core.image.Image(image, mask=mask), workspace
            )
            assert round(abs(l1 - l), 7) == 0
//  [python]
def test_small_images():
    """Test mixture of gaussians thresholding with few pixels

    Run MOG to see if it blows up, given 0-10 pixels"""
    r = numpy.random.RandomState()
    r.seed(91)
    image = r.uniform(size=(9, 11))
    ii, jj = numpy.mgrid[0 : image.shape[0], 0 : image.shape[1]]
    ii, jj = ii.flatten(), jj.flatten()

    for threshold_method in (
        cellprofiler.modules.threshold.TM_LI,
        cellprofiler.modules.threshold.TM_OTSU,
        cellprofiler.modules.threshold.TM_ROBUST_BACKGROUND,
    ):
        for i in range(11):
            mask = numpy.zeros(image.shape, bool)
            if i:
                p = r.permutation(numpy.prod(image.shape))[:i]
                mask[ii[p], jj[p]] = True
            workspace, x = make_workspace(image, mask)
            x.global_operation.value = threshold_method
            x.threshold_scope.value = TS_GLOBAL
            l, g, _, _, _ = x.get_threshold(
                cellprofiler_core.image.Image(image, mask=mask), workspace
            )
            v = image[mask]
            image = r.uniform(size=(9, 11))
            image[mask] = v
            l1, g1, _, _, _ = x.get_threshold(
                cellprofiler_core.image.Image(image, mask=mask), workspace
            )
            assert round(abs(l1 - l), 7) == 0
//  [python]
def test_initial_list(self, initial):
        expected_exc = ValueError
        arr = [[1, 2], [3, 4]]
        # Numpy raises ValueError:
        # Input object to FillWithScalar is not a scalar
        with pytest.raises(expected_exc):
            np.prod(arr, initial=initial)
        # when LEGATE_TEST=1, cuNumeric casts list to scalar and proceeds
        with pytest.raises(expected_exc):
            num.prod(arr, initial=initial)
//  [python]
def hessian_to_block_tree(hessian, f_tree, params_tree):
    """Convert a Hessian array to block-tree format.

    Remark: In comparison to Jax we need this formatting function because we calculate
    the second derivative using second-order finite differences. Jax computes the
    second derivative by applying their jacobian function twice, which produces the
    desired block-tree shape of the Hessian automatically. If we apply our first
    derivative function twice we get the same block-tree shape.

    Args:
        hessian (np.ndarray): The Hessian, 2- or 3-dimensional array representation of
            the resulting block-tree.
        f_tree (pytree): The function evaluated at params_tree.
        params_tree (pytree): The params_tree.

    Returns:
        hessian_block_tree (pytree): The pytree

    """
    _check_dimensions_hessian(hessian, f_tree, params_tree)

    if hessian.ndim == 2:
        hessian = hessian[np.newaxis]

    flat_f, treedef_f = tree_flatten(f_tree)
    flat_p, treedef_p = tree_flatten(params_tree)

    flat_f_np = [_convert_to_numpy(leaf, only_pandas=True) for leaf in flat_f]
    flat_p_np = [_convert_to_numpy(leaf, only_pandas=True) for leaf in flat_p]

    shapes_f = [np.shape(a) for a in flat_f_np]
    shapes_p = [np.shape(a) for a in flat_p_np]

    block_bounds_f = np.cumsum([int(np.prod(s)) for s in shapes_f[:-1]])
    block_bounds_p = np.cumsum([int(np.prod(s)) for s in shapes_p[:-1]])

    sub_block_trees = []
    for s0, subarr in zip(shapes_f, np.split(hessian, block_bounds_f, axis=0)):
        blocks = []
        for leaf_outer, s1, submat in zip(
            flat_p, shapes_p, np.split(subarr, block_bounds_p, axis=1)
        ):
            row = []
            for leaf_inner, s2, block_values in zip(
                flat_p, shapes_p, np.split(submat, block_bounds_p, axis=2)
            ):
                _shape = [k for k in (*s0, *s1, *s2) if k != 1]
                raw_block = block_values.reshape(_shape)
                block = _convert_raw_block_to_pandas(raw_block, leaf_outer, leaf_inner)
                row.append(block)
            blocks.append(row)
        block_tree = tree_unflatten(
            treedef_p, [tree_unflatten(treedef_p, row) for row in blocks]
        )
        sub_block_trees.append(block_tree)

    hessian_block_tree = tree_unflatten(treedef_f, sub_block_trees)
    return hessian_block_tree
//  [python]
def test_initial_empty_array(self):
        size = (1, 0)
        arr_np = np.random.random(size) * 10
        arr_num = num.array(arr_np)
        initial_value = np.random.uniform(-20.0, 20.0)
        out_num = num.prod(arr_num, initial=initial_value)
        out_np = np.prod(arr_np, initial=initial_value)
        assert allclose(out_np, out_num)
//  [python]
def test_basic(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        out_np = np.prod(arr_np)
        out_num = arr_num.prod()
        assert allclose(out_np, out_num)
        assert allclose(out_num, arr_num.prod())
//  [python]
def test_basic(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        out_np = np.prod(arr_np)
        out_num = np.prod(arr_num)
        assert allclose(out_np, out_num)
        assert allclose(out_num, arr_num.prod())
//  [python]
def test_where(self):
        arr = [[1, 2], [3, 4]]
        out_np = np.prod(arr, where=[False, True])
        out_num = num.prod(arr, where=[False, True])
        assert allclose(out_np, out_num)

        # where is boolean
        out_np = np.prod(arr, where=True)
        out_num = num.prod(arr, where=True)
        assert allclose(out_np, out_num)

        out_np = np.prod(arr, where=False)
        out_num = num.prod(arr, where=False)
        assert allclose(out_np, out_num)
//  [python]
def test_dtype(self, dtype):
        size = (5, 5, 5)
        arr = np.random.random(size) * 10
        arr_np = np.array(arr, dtype=dtype)
        arr_num = num.array(arr_np)
        out_np = np.prod(arr_np)
        out_num = num.prod(arr_num)
        assert allclose(out_np, out_num)
//  [python]
def test_dtype_integer_precision(self, dtype):
        arr_np = np.arange(0, 5).astype(dtype)
        arr_num = num.arange(0, 5).astype(dtype)
        out_np = np.prod(arr_np)
        out_num = num.prod(arr_num)
        assert allclose(out_num, arr_num.prod())
        # When input precision is less than default platform integer
        # NumPy returns the product with dtype of platform integer
        # cuNumeric returns the product with dtype of the input array
        assert allclose(out_np, out_num)
//  [python]
def test_axis_basic(self, axis):
        size = (5, 5, 5)
        arr_np = np.random.random(size) * 10
        arr_num = num.array(arr_np)
        out_num = num.prod(arr_num, axis=axis)
        out_np = np.prod(arr_np, axis=axis)
        assert allclose(out_np, out_num)
//  [python]
def test_dtype_complex(self, dtype):
        arr = (np.random.rand(5, 5) * 10 + 2) + (
            np.random.rand(5, 5) * 10 * 1.0j + 0.2j
        )
        arr_np = np.array(arr, dtype=dtype)
        arr_num = num.array(arr, dtype=dtype)
        out_np = np.prod(arr_np)
        # cunumeric always returns [1+0.j] when LEGATE_TEST=1
        out_num = num.prod(arr_num)
        # When running tests with CUNUMERIC_TEST=1 and dtype is complex256,
        # allclose hits assertion error:
        # File "/legate/cunumeric/cunumeric/eager.py", line 293,
        # in to_deferred_array
        #   assert self.runtime.is_supported_type(self.array.dtype)
        #   AssertionError
        assert allclose(out_np, out_num)
//  [python]
def test_out_basic(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        out_np = np.random.random(())
        out_num = num.random.random(())
        np.multiply.reduce(arr_np, out=out_np)
        num.prod(arr_num, out=out_num)
        assert allclose(out_np, out_num)
//  [python]
def test_out_basic(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        out_np = np.random.random(())
        out_num = num.random.random(())
        np.prod(arr_np, out=out_np)
        num.prod(arr_num, out=out_num)
        assert allclose(out_np, out_num)
//  [python]
def test_axis_tuple(self, axis):
        size = (5, 5, 5)
        arr_np = np.random.random(size) * 10
        arr_num = num.array(arr_np)
        out_np = np.prod(arr_np, axis=axis)
        # when LEGATE_TEST = 1 cuNumeric raises two types of exceptions
        # (-1, 1): ValueError: Invalid promotion on dimension 2 for a 1-D store
        # others:
        # NotImplementedError: Need support for reducing multiple dimensions
        # Numpy get results.
        out_num = num.prod(arr_num, axis=axis)
        assert allclose(out_np, out_num)
//  [python]
def test_axis_keepdims_false(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        ndim = arr_np.ndim
        for axis in range(-ndim + 1, ndim, 1):
            out_np = np.prod(arr_np, axis=axis, keepdims=False)
            out_num = num.prod(arr_num, axis=axis, keepdims=False)
            assert allclose(out_np, out_num)
//  [python]
def test_out_axis(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        ndim = arr_np.ndim
        for axis in range(-ndim + 1, ndim, 1):
            out_shape = ()
            if isinstance(size, tuple):
                out_shape_list = list(size)
                del out_shape_list[axis]
                out_shape = tuple(out_shape_list)
            out_np = np.random.random(out_shape)
            out_num = num.random.random(out_shape)
            np.prod(arr_np, out=out_np, axis=axis)
            num.prod(arr_num, out=out_num, axis=axis)
            assert allclose(out_np, out_num)
//  [python]
def numel(x):
    '''
    numel(x) yields the number of elements in x: the product of the shape of x.
    '''
    return int(np.prod(np.shape(x)))
//  [python]
def test_initial(self, size):
        arr_np = np.random.random(size) * 10
        arr_num = num.array(arr_np)
        initial_value = np.random.uniform(-20.0, 20.0)
        out_num = num.prod(arr_num, initial=initial_value)
        out_np = np.prod(arr_np, initial=initial_value)

        assert allclose(out_np, out_num)
//  [python]
def test_out_axis_dtype(self, size):
        arr = np.random.random(size) * 10
        arr_np = np.array(arr, dtype=to_dtype("f"))
        arr_num = num.array(arr, dtype=to_dtype("f"))

        ndim = arr_np.ndim
        for axis in range(-ndim + 1, ndim, 1):
            out_shape = ()
            if isinstance(size, tuple):
                out_shape_list = list(size)
                del out_shape_list[axis]
                out_shape = tuple(out_shape_list)
            out = np.random.random(out_shape)

            out_np = np.array(out, dtype=to_dtype("i"))
            out_num = num.array(out, dtype=to_dtype("i"))

            np.multiply.reduce(arr_np, out=out_np, axis=axis)
            num.multiply.reduce(arr_num, out=out_num, axis=axis)

            assert allclose(out_np, out_num)
//  [python]
def test_out_axis_dtype(self, size):
        arr = np.random.random(size) * 10
        arr_np = np.array(arr, dtype=to_dtype("f"))
        arr_num = num.array(arr, dtype=to_dtype("f"))

        ndim = arr_np.ndim
        for axis in range(-ndim + 1, ndim, 1):
            out_shape = ()
            if isinstance(size, tuple):
                out_shape_list = list(size)
                del out_shape_list[axis]
                out_shape = tuple(out_shape_list)
            out = np.random.random(out_shape)

            out_np = np.array(out, dtype=to_dtype("i"))
            out_num = num.array(out, dtype=to_dtype("i"))

            np.prod(arr_np, out=out_np, axis=axis)
            num.prod(arr_num, out=out_num, axis=axis)

            assert allclose(out_np, out_num)
//  [python]
def test_axis_keepdims_true(self, size):
        arr_np = np.random.random(size)
        arr_num = num.array(arr_np)
        ndim = arr_np.ndim
        for axis in range(-ndim + 1, ndim, 1):
            out_np = np.prod(arr_np, axis=axis, keepdims=True)
            out_num = num.prod(arr_num, axis=axis, keepdims=True)
            # in cunumeric/deferred/unary_reduction:
            # if lhs_array.size == 1:
            #     > assert axes is None or len(axes) == rhs_array.ndim - (
            #         0 if keepdims else lhs_array.ndim
            #     )
            # E    AssertionError
            assert allclose(out_np, out_num)
//  [python]
def number_of_values(self):
        return n.prod(self.dimensions)
//  [python]
def number_of_values(self):
        return numpy.prod(self.dimensions)
//  [python]
def _compute_multiplier(limit: int, dtype, largest_block: int, result):
    """
    Utility function for auto_chunk, to fin how much larger or smaller the ideal
    chunk size is relative to what we have now.
    """
    return (
        limit
        / dtype.itemsize
        / largest_block
        / numpy.prod(list(r if r != 0 else 1 for r in result.values()))
    )
//  [python]
def _compute_multiplier(limit: int, dtype, largest_block: int, result):
    """
    Utility function for auto_chunk, to fin how much larger or smaller the ideal
    chunk size is relative to what we have now.
    """
    return (
        limit
        / dtype.itemsize
        / largest_block
        / np.prod(list(r if r != 0 else 1 for r in result.values()))
    )
//  [python]
def rand(*args,**kwargs):
    if not args:
        return np.random.rand()
    if len(args) == 1:
        args += args
    try:
        return np.random.rand(np.prod(args)).reshape(args,order="F")
    except:
        pass
//  [python]
def _repr_html_table(self):
        if "sparse" in typename(type(self._meta)):
            nbytes = None
            cbytes = None
        elif not math.isnan(self.nbytes):
            nbytes = format_bytes(self.nbytes)
            cbytes = format_bytes(n.prod(self.chunksize) * self.dtype.itemsize)
        else:
            nbytes = "unknown"
            cbytes = "unknown"

        table = [
            "<table>",
            "  <thead>",
            "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>",
            "  </thead>",
            "  <tbody>",
            "    <tr><th> Bytes </th><td> %s </td> <td> %s </td></tr>"
            % (nbytes, cbytes)
            if nbytes is not None
            else "",
            "    <tr><th> Shape </th><td> %s </td> <td> %s </td></tr>"
            % (str(self.shape), str(self.chunksize)),
            "    <tr><th> Count </th><td> %d Tasks </td><td> %d Chunks </td></tr>"
            % (len(self.__dask_graph__()), self.npartitions),
            "    <tr><th> Type </th><td> %s </td><td> %s.%s </td></tr>"
            % (
                self.dtype,
                type(self._meta).__module__.split(".")[0],
                type(self._meta).__name__,
            ),
            "  </tbody>",
            "</table>",
        ]
        return "\n".join(table)
//  [python]
def _repr_html_table(self):
        if "sparse" in typename(type(self._meta)):
            nbytes = None
            cbytes = None
        elif not math.isnan(self.nbytes):
            nbytes = format_bytes(self.nbytes)
            cbytes = format_bytes(np.prod(self.chunksize) * self.dtype.itemsize)
        else:
            nbytes = "unknown"
            cbytes = "unknown"

        table = [
            "<table>",
            "  <thead>",
            "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>",
            "  </thead>",
            "  <tbody>",
            "    <tr><th> Bytes </th><td> %s </td> <td> %s </td></tr>"
            % (nbytes, cbytes)
            if nbytes is not None
            else "",
            "    <tr><th> Shape </th><td> %s </td> <td> %s </td></tr>"
            % (str(self.shape), str(self.chunksize)),
            "    <tr><th> Count </th><td> %d Tasks </td><td> %d Chunks </td></tr>"
            % (len(self.__dask_graph__()), self.npartitions),
            "    <tr><th> Type </th><td> %s </td><td> %s.%s </td></tr>"
            % (
                self.dtype,
                type(self._meta).__module__.split(".")[0],
                type(self._meta).__name__,
            ),
            "  </tbody>",
            "</table>",
        ]
        return "\n".join(table)
//  [python]
def randn(*args,**kwargs):
    if not args:
        return np.random.randn()
    if len(args) == 1:
        args += args
    try:
        return np.random.randn(np.prod(args)).reshape(args,order="F")
    except:
        pass
//  [python]
def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
  a = asarray(a).data
  original_shape = a._shape_as_list()  # pylint: disable=protected-access
  # Best effort recovery of the shape.
  if original_shape is not None and None not in original_shape:
    if not original_shape:
      original_shape = (repeats,)
    else:
      repeats_np = np.ravel(np.array(repeats))
      if repeats_np.size == 1:
        repeats_np = repeats_np.item()
        if axis is None:
          original_shape = (repeats_np * n.prod(original_shape),)
        else:
          original_shape[axis] = repeats_np * original_shape[axis]
      else:
        if axis is None:
          original_shape = (repeats_np.sum(),)
        else:
          original_shape[axis] = repeats_np.sum()

  repeats = asarray(repeats).data
  result = tf.repeat(a, repeats, axis)
  result.set_shape(original_shape)

  return utils.tensor_to_ndarray(result)
//  [python]
def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
  a = asarray(a).data
  original_shape = a._shape_as_list()  # pylint: disable=protected-access
  # Best effort recovery of the shape.
  if original_shape is not None and None not in original_shape:
    if not original_shape:
      original_shape = (repeats,)
    else:
      repeats_np = np.ravel(np.array(repeats))
      if repeats_np.size == 1:
        repeats_np = repeats_np.item()
        if axis is None:
          original_shape = (repeats_np * np.prod(original_shape),)
        else:
          original_shape[axis] = repeats_np * original_shape[axis]
      else:
        if axis is None:
          original_shape = (repeats_np.sum(),)
        else:
          original_shape[axis] = repeats_np.sum()

  repeats = asarray(repeats).data
  result = tf.repeat(a, repeats, axis)
  result.set_shape(original_shape)

  return utils.tensor_to_ndarray(result)
//  [python]
def accuracy(pred, label, topk=(1,)):
    maxk = max(topk)
    pred = np.argsort(pred)[..., ::-1][..., :maxk]
    if len(label.shape) == 1:
        label = label.reshape(-1, 1)
    elif label.shape[-1] != 1:
        label = np.argmax(label, axis=-1)
        label = label[..., np.newaxis]
    correct = pred == np.repeat(label, maxk, -1)

    total = np.multiply.reduce(np.array(label.shape[:-1]))

    res = []
    for k in topk:
        correct_k = correct[..., :k].sum()
        res.append(float(correct_k) / total)
    return res
//  [python]
def accuracy(pred, label, topk=(1,)):
    maxk = max(topk)
    pred = np.argsort(pred)[..., ::-1][..., :maxk]
    if len(label.shape) == 1:
        label = label.reshape(-1, 1)
    elif label.shape[-1] != 1:
        label = np.argmax(label, axis=-1)
        label = label[..., np.newaxis]
    correct = pred == np.repeat(label, maxk, -1)

    total = np.prod(np.array(label.shape[:-1]))

    res = []
    for k in topk:
        correct_k = correct[..., :k].sum()
        res.append(float(correct_k) / total)
    return res
//  [python]
def unbroadcast(a, b):
    '''
    unbroadcast(a, b) yields a tuple (aa, bb) that is equivalent to (a, b) except that aa and bb
      have been reshaped such that arithmetic numpy operations such as aa * bb will result in
      row-wise operation instead of column-wise broadcasting.
    '''
    # First: promote the arrays.
    (a,b) = [asdense(x) for x in promote(a, b)]
    # See if they're tensors first
    try:
        torch = pytorch()
        if torch.is_tensor(a):
            da = len(a.shape)
            db = len(b.shape)
            if   da > db: return (a, torch.reshape(b, b.shape + (1,)*(da - db)))
            elif da < db: return (torch.reshape(a, a.shape + (1,)*(db - da)), b)
            else:         return (a, b)
    except ImportError: pass
    # they could be sparse:
    spa = sps.issparse(a)
    spb = sps.issparse(b)
    if   spa and spb: return (a,b)
    elif spa or  spb:
        def fix(sp,nm):
            nm = np.asarray(nm)
            dnm = len(nm.shape)
            nnm = np.multiply.reduce(nm.shape)
            # if we have (sparse matrix) * (high-dim array), unbroadcast the dense array
            if   dnm == 0: return (sp, np.reshape(nm, (1,   1)))
            elif dnm == 1: return (sp, np.reshape(nm, (nnm, 1)))
            elif dnm == 2: return (sp, nm)
            else:          return unbroadcast(sp.toarray(), nm)
        return fix(a, b) if spa else tuple(reversed(fix(b, a)))
    # okay, no sparse matrices found:
    da = len(a.shape)
    db = len(b.shape)
    if   da > db: return (a, np.reshape(b, b.shape + (1,)*(da - db)))
    elif da < db: return (np.reshape(a, a.shape + (1,)*(db - da)), b)
    else:         return (a, b)
//  [python]
def unbroadcast(a, b):
    '''
    unbroadcast(a, b) yields a tuple (aa, bb) that is equivalent to (a, b) except that aa and bb
      have been reshaped such that arithmetic numpy operations such as aa * bb will result in
      row-wise operation instead of column-wise broadcasting.
    '''
    # First: promote the arrays.
    (a,b) = [asdense(x) for x in promote(a, b)]
    # See if they're tensors first
    try:
        torch = pytorch()
        if torch.is_tensor(a):
            da = len(a.shape)
            db = len(b.shape)
            if   da > db: return (a, torch.reshape(b, b.shape + (1,)*(da - db)))
            elif da < db: return (torch.reshape(a, a.shape + (1,)*(db - da)), b)
            else:         return (a, b)
    except ImportError: pass
    # they could be sparse:
    spa = sps.issparse(a)
    spb = sps.issparse(b)
    if   spa and spb: return (a,b)
    elif spa or  spb:
        def fix(sp,nm):
            nm = np.asarray(nm)
            dnm = len(nm.shape)
            nnm = np.prod(nm.shape)
            # if we have (sparse matrix) * (high-dim array), unbroadcast the dense array
            if   dnm == 0: return (sp, np.reshape(nm, (1,   1)))
            elif dnm == 1: return (sp, np.reshape(nm, (nnm, 1)))
            elif dnm == 2: return (sp, nm)
            else:          return unbroadcast(sp.toarray(), nm)
        return fix(a, b) if spa else tuple(reversed(fix(b, a)))
    # okay, no sparse matrices found:
    da = len(a.shape)
    db = len(b.shape)
    if   da > db: return (a, np.reshape(b, b.shape + (1,)*(da - db)))
    elif da < db: return (np.reshape(a, a.shape + (1,)*(db - da)), b)
    else:         return (a, b)
//  [python]
def stack_images( images ):
    images_shape = numpy.array( images.shape )
    new_axes = get_transpose_axes( len( images_shape ) )
    new_shape = [ numpy.multiply.reduce( images_shape[x] ) for x in new_axes ]
    return numpy.transpose(
        images,
        axes = numpy.concatenate( new_axes )
        ).reshape( new_shape )
//  [python]
def stack_images( images ):
    images_shape = numpy.array( images.shape )
    new_axes = get_transpose_axes( len( images_shape ) )
    new_shape = [ numpy.prod( images_shape[x] ) for x in new_axes ]
    return numpy.transpose(
        images,
        axes = numpy.concatenate( new_axes )
        ).reshape( new_shape )