jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_ptp(self):
        x = np.arange(4).reshape((2, 2))
        mx = x.view(np.matrix)
        assert_(mx.ptp() == 3)
        assert_(mx.ptp(0).all() == np.array([2, 2]).all())
        assert_(mx.ptp(1).all() == np.array([1, 1]).all())
//  [python]
def test_ptp(self):
        x = np.arange(4).reshape((2, 2))
        mx = x.view(np.matrix)
        assert_(mx.ptp() == 3)
        assert_(np.all(mx.ptp(0) == np.array([2, 2])))
        assert_(np.all(mx.ptp(1) == np.array([1, 1])))
//  [python]
def bound_perturbed_gmres(pseudo, p, epsilon, deltas):
    """Compute GMRES perturbation bound based on pseudospectrum

    Computes the GMRES bound from [SifEM13]_.
    """
    if not n.all(n.array(deltas) > epsilon):
        raise ArgumentError("all deltas have to be greater than epsilon")

    bound = []
    for delta in deltas:
        # get boundary paths
        paths = pseudo.contour_paths(delta)

        # get vertices on boundary
        vertices = paths.vertices()

        # evaluate polynomial
        supremum = numpy.max(numpy.abs(p(vertices)))

        # compute bound
        bound.append(
            epsilon
            / (delta - epsilon)
            * paths.length()
            / (2 * numpy.pi * delta)
            * supremum
        )
    return bound
//  [python]
def bound_perturbed_gmres(pseudo, p, epsilon, deltas):
    """Compute GMRES perturbation bound based on pseudospectrum

    Computes the GMRES bound from [SifEM13]_.
    """
    if not numpy.all(numpy.array(deltas) > epsilon):
        raise ArgumentError("all deltas have to be greater than epsilon")

    bound = []
    for delta in deltas:
        # get boundary paths
        paths = pseudo.contour_paths(delta)

        # get vertices on boundary
        vertices = paths.vertices()

        # evaluate polynomial
        supremum = numpy.max(numpy.abs(p(vertices)))

        # compute bound
        bound.append(
            epsilon
            / (delta - epsilon)
            * paths.length()
            / (2 * numpy.pi * delta)
            * supremum
        )
    return bound
//  [python]
def __getitem__(self, idx):
        """ Slicing data from the array.

        Parameters
        ----------
        idx : Numpy array or Pytorch tensor or slice or int.
            The index.

        Returns
        -------
        Numpy array : the data from the HDF5 array indexed by `idx`.
        """
        if isinstance(idx, (slice, int)):
            return self._arr[idx].astype(self._out_dtype)

        if isinstance(idx, th.Tensor):
            idx = idx.numpy()
        # If the idx are sorted.
        if np.all(idx[1:] - idx[:-1] > 0):
            arr = self._arr[idx]
        else:
            # There are two cases here: 1) there are duplicated IDs,
            # 2) the IDs are not sorted. Unique can return unique
            # IDs in the ascending order that meets the requirement of
            # HDF5 indexing.
            uniq_ids, reverse_idx = np.unique(idx, return_inverse=True)
            arr = self._arr[uniq_ids][reverse_idx]

        if self._out_dtype is not None:
            arr = arr.astype(self._out_dtype)
        return arr
//  [python]
def new_structured_grid(shape, cell_corners, cell_positions=None, axes_names = "xyz", offset=None):
        """Returns a structured grid with cells with given corners and cell_positions.
           if not present, cell positions default to average of corner positions.
        """
        if len(axes_names)<len(shape):
            raise Exception("provide enough axes names")
        if len(cell_corners)!=len(shape):
            raise Exception("dimensions of shape and cell_boundaries do not conform")
        for c in cell_corners:
            if not n.all([s1==s2+1 for s1,s2 in zip(c.shape,shape)]):
                shape1=[s+1 for s in shape]
                raise Exception("size of cell_corner arrays must be {0} instead of {1}".format(shape1.__str__(),c.shape.__str__()))        

        if cell_positions is None:
              cell_positions=[]
              for cc in cell_corners:
                  cp=n.zeros(shape) * cc.flat[0]
                  for i in range(2**len(shape)):
                      slicing=()
                      for j in range(len(shape)):
                          if i & 2**j:
                              slicing+=(slice(1,None),) 
                          else:
                              slicing+=(slice(None,-1),)
                      cp=cp+cc[slicing]              
                  cell_positions.append(cp/2**len(shape))          

        if len(cell_positions)!=len(shape):
            raise Exception("dimensions of shape and cell_positions do not conform")
        for c in cell_positions:
            if not n.all([s1==s2 for s1,s2 in zip(c.shape,shape)]):
                raise Exception("size of cell_position arrays must be {0} instead of {1}".format(shape1.__str__(),c.shape.__str__()))        

        if offset is None:
            offset=[0.*l.flat[0] for l in cell_positions]

        result = StructuredGrid(*shape)

        for axis_name, pos, of in zip(axes_names, cell_positions, offset):
            setattr(result, axis_name, pos + of)

        result.add_vector_attribute("position", axes_names[0:len(shape)])
        
        object.__setattr__(result,"_grid_type","structured")
        object.__setattr__(result,"_cell_corners", cell_corners)
       
        return result
//  [python]
def new_structured_grid(shape, cell_corners, cell_positions=None, axes_names = "xyz", offset=None):
        """Returns a structured grid with cells with given corners and cell_positions.
           if not present, cell positions default to average of corner positions.
        """
        if len(axes_names)<len(shape):
            raise Exception("provide enough axes names")
        if len(cell_corners)!=len(shape):
            raise Exception("dimensions of shape and cell_boundaries do not conform")
        for c in cell_corners:
            if not numpy.all([s1==s2+1 for s1,s2 in zip(c.shape,shape)]):
                shape1=[s+1 for s in shape]
                raise Exception("size of cell_corner arrays must be {0} instead of {1}".format(shape1.__str__(),c.shape.__str__()))        

        if cell_positions is None:
              cell_positions=[]
              for cc in cell_corners:
                  cp=numpy.zeros(shape) * cc.flat[0]
                  for i in range(2**len(shape)):
                      slicing=()
                      for j in range(len(shape)):
                          if i & 2**j:
                              slicing+=(slice(1,None),) 
                          else:
                              slicing+=(slice(None,-1),)
                      cp=cp+cc[slicing]              
                  cell_positions.append(cp/2**len(shape))          

        if len(cell_positions)!=len(shape):
            raise Exception("dimensions of shape and cell_positions do not conform")
        for c in cell_positions:
            if not numpy.all([s1==s2 for s1,s2 in zip(c.shape,shape)]):
                raise Exception("size of cell_position arrays must be {0} instead of {1}".format(shape1.__str__(),c.shape.__str__()))        

        if offset is None:
            offset=[0.*l.flat[0] for l in cell_positions]

        result = StructuredGrid(*shape)

        for axis_name, pos, of in zip(axes_names, cell_positions, offset):
            setattr(result, axis_name, pos + of)

        result.add_vector_attribute("position", axes_names[0:len(shape)])
        
        object.__setattr__(result,"_grid_type","structured")
        object.__setattr__(result,"_cell_corners", cell_corners)
       
        return result
//  [python]
def test_identity():
    def check(dtype):
        obj = random_of_dtype((2,), dtype)
        sym = vector(dtype=dtype)
        f = function([sym], tensor_copy(sym))
        assert n.all(obj == f(obj))
        assert obj.dtype == f(obj).dtype
        topo = f.maker.fgraph.toposort()
        assert len(topo) == 1
        if config.mode != "FAST_COMPILE":
            assert isinstance(topo[0].op, DeepCopyOp)

    for dtype in ALL_DTYPES:
        check(dtype)
//  [python]
def test_identity():
    def check(dtype):
        obj = random_of_dtype((2,), dtype)
        sym = vector(dtype=dtype)
        f = function([sym], tensor_copy(sym))
        assert np.all(obj == f(obj))
        assert obj.dtype == f(obj).dtype
        topo = f.maker.fgraph.toposort()
        assert len(topo) == 1
        if config.mode != "FAST_COMPILE":
            assert isinstance(topo[0].op, DeepCopyOp)

    for dtype in ALL_DTYPES:
        check(dtype)
//  [python]
def test_stack_scalar_make_vector(self):
        # Test that calling stack() on scalars instantiates MakeVector,
        # not Join. Test that the floatX dtype stay floatX, not downcasted
        # to int64
        a = scalar("a", dtype=self.floatX)
        b = scalar("b", dtype=self.floatX)
        s = stack([a, b, a, b])
        f = function([a, b], s, mode=self.mode)
        val = f(1, 2)
        # print val
        assert np.all(val == [1, 2, 1, 2])
        topo = f.maker.fgraph.toposort()
        assert len([n for n in topo if isinstance(n.op, MakeVector)]) > 0
        assert len([n for n in topo if isinstance(n, type(self.join_op))]) == 0
        assert f.maker.fgraph.outputs[0].dtype == self.floatX
//  [python]
def test_stack_scalar_make_vector_dtype(self):
        # Test that calling stack() on scalars instantiates MakeVector,
        # event when the scalar don't have the same dtype.
        a = iscalar("a")
        b = lscalar("b")
        s = stack([a, b, a, b])
        f = function([a, b], s, mode=self.mode)
        val = f(1, 2)
        assert n.all(val == [1, 2, 1, 2])
        topo = f.maker.fgraph.toposort()
        assert len([n for n in topo if isinstance(n.op, MakeVector)]) > 0
        assert len([n for n in topo if isinstance(n, type(self.join_op))]) == 0
        assert f.maker.fgraph.outputs[0].dtype == "int64"
//  [python]
def test_stack_scalar_make_vector_dtype(self):
        # Test that calling stack() on scalars instantiates MakeVector,
        # event when the scalar don't have the same dtype.
        a = iscalar("a")
        b = lscalar("b")
        s = stack([a, b, a, b])
        f = function([a, b], s, mode=self.mode)
        val = f(1, 2)
        assert np.all(val == [1, 2, 1, 2])
        topo = f.maker.fgraph.toposort()
        assert len([n for n in topo if isinstance(n.op, MakeVector)]) > 0
        assert len([n for n in topo if isinstance(n, type(self.join_op))]) == 0
        assert f.maker.fgraph.outputs[0].dtype == "int64"
//  [python]
def test_op_cache():
    # TODO: What is this actually testing?
    # trigger bug in ticket #162
    v = matrix()
    v.name = "v"
    gv = fill(v / v, 1.0) / v - (fill(v / v, 1.0) * v) / (v * v)
    fn_py = inplace_func([v], gv)
    fn_c_or_py = inplace_func([v], gv)

    a = random(5, 2).astype(config.floatX)
    assert n.all(fn_py(a) == fn_c_or_py(a))
//  [python]
def test_op_cache():
    # TODO: What is this actually testing?
    # trigger bug in ticket #162
    v = matrix()
    v.name = "v"
    gv = fill(v / v, 1.0) / v - (fill(v / v, 1.0) * v) / (v * v)
    fn_py = inplace_func([v], gv)
    fn_c_or_py = inplace_func([v], gv)

    a = random(5, 2).astype(config.floatX)
    assert np.all(fn_py(a) == fn_c_or_py(a))
//  [python]
def test_stack_scalar_make_vector_constant(self):
        # Test that calling stack() on scalars instantiates MakeVector,
        # event when the scalar are simple int type.
        a = iscalar("a")
        b = lscalar("b")
        # test when the constant is the first element.
        # The first element is used in a special way
        s = stack([10, a, b, np.int8(3)])
        f = function([a, b], s, mode=self.mode)
        val = f(1, 2)
        assert n.all(val == [10, 1, 2, 3])
        topo = f.maker.fgraph.toposort()
        assert len([n for n in topo if isinstance(n.op, MakeVector)]) > 0
        assert len([n for n in topo if isinstance(n, type(self.join_op))]) == 0
        assert f.maker.fgraph.outputs[0].dtype == "int64"
//  [python]
def test_stack_scalar_make_vector_constant(self):
        # Test that calling stack() on scalars instantiates MakeVector,
        # event when the scalar are simple int type.
        a = iscalar("a")
        b = lscalar("b")
        # test when the constant is the first element.
        # The first element is used in a special way
        s = stack([10, a, b, np.int8(3)])
        f = function([a, b], s, mode=self.mode)
        val = f(1, 2)
        assert np.all(val == [10, 1, 2, 3])
        topo = f.maker.fgraph.toposort()
        assert len([n for n in topo if isinstance(n.op, MakeVector)]) > 0
        assert len([n for n in topo if isinstance(n, type(self.join_op))]) == 0
        assert f.maker.fgraph.outputs[0].dtype == "int64"
//  [python]
def test_flatten_ndim_default():
    a = dmatrix()
    c = flatten(a)
    f = inplace_func([a], c)
    a_val = _asarray([[0, 1, 2], [3, 4, 5]], dtype="float64")
    c_val = _asarray([0, 1, 2, 3, 4, 5], dtype="float64")
    assert n.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert n.all(f(a_val) == c_val)

    utt.verify_grad(flatten, [a_val])
//  [python]
def test_flatten_ndim_default():
    a = dmatrix()
    c = flatten(a)
    f = inplace_func([a], c)
    a_val = _asarray([[0, 1, 2], [3, 4, 5]], dtype="float64")
    c_val = _asarray([0, 1, 2, 3, 4, 5], dtype="float64")
    assert np.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert np.all(f(a_val) == c_val)

    utt.verify_grad(flatten, [a_val])
//  [python]
def test_flatten_scalar():
    a = dscalar()
    c = flatten(a)
    f = inplace_func([a], c)
    a_val = _asarray(3.0, dtype="float64")
    c_val = _asarray([3.0], dtype="float64")
    assert n.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert n.all(f(a_val) == c_val)
//  [python]
def test_flatten_scalar():
    a = dscalar()
    c = flatten(a)
    f = inplace_func([a], c)
    a_val = _asarray(3.0, dtype="float64")
    c_val = _asarray([3.0], dtype="float64")
    assert np.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert np.all(f(a_val) == c_val)
//  [python]
def test_flatten_ndim1():
    a = dmatrix()
    c = flatten(a, 1)
    f = inplace_func([a], c)
    a_val = _asarray([[0, 1, 2], [3, 4, 5]], dtype="float64")
    c_val = _asarray([0, 1, 2, 3, 4, 5], dtype="float64")
    assert n.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert n.all(f(a_val) == c_val)

    utt.verify_grad(flatten, [a_val])
//  [python]
def test_flatten_ndim1():
    a = dmatrix()
    c = flatten(a, 1)
    f = inplace_func([a], c)
    a_val = _asarray([[0, 1, 2], [3, 4, 5]], dtype="float64")
    c_val = _asarray([0, 1, 2, 3, 4, 5], dtype="float64")
    assert np.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert np.all(f(a_val) == c_val)

    utt.verify_grad(flatten, [a_val])
//  [python]
def test_flatten_ndim2():
    a = dmatrix()
    c = flatten(a, 2)
    f = inplace_func([a], c)
    a_val = _asarray([[0, 1, 2], [3, 4, 5]], dtype="float64")
    assert n.all(f(a_val) == a_val)
    f = inplace_func([a], c)
    assert n.all(f(a_val) == a_val)

    flatten_2 = partial(flatten, ndim=2)
    utt.verify_grad(flatten_2, [a_val])
//  [python]
def test_flatten_ndim2():
    a = dmatrix()
    c = flatten(a, 2)
    f = inplace_func([a], c)
    a_val = _asarray([[0, 1, 2], [3, 4, 5]], dtype="float64")
    assert np.all(f(a_val) == a_val)
    f = inplace_func([a], c)
    assert np.all(f(a_val) == a_val)

    flatten_2 = partial(flatten, ndim=2)
    utt.verify_grad(flatten_2, [a_val])
//  [python]
def test_stack_new_interface(self):
        # Test the new numpy-like interface: stack(tensors, axis=0).

        a = imatrix("a")
        b = imatrix("b")

        # Testing axis parameter
        s3 = stack([a, b], 1)
        f = function([a, b], s3, mode=self.mode)
        v3 = f([[1, 2]], [[3, 4]])
        v4 = np.array([[[1, 2], [3, 4]]])
        assert v3.shape == v4.shape
        assert n.all(v3 == v4)
        # Testing negative axis
        v1 = [[1, 2, 3], [4, 5, 6]]
        v2 = [[7, 8, 9], [10, 11, 12]]
        s = stack([a, b], axis=-1)
        f = function([a, b], s, mode=self.mode)
        v = np.zeros((2, 3, 2))
        v[:, :, 0] = v1
        v[:, :, 1] = v2
        out = f(v1, v2)
        assert v.shape == out.shape
        assert n.all(v == out)
        s = stack([a, b], axis=-2)
        f = function([a, b], s, mode=self.mode)
        v = np.zeros((2, 2, 3))
        v[:, 0, :] = v1
        v[:, 1, :] = v2
        out = f(v1, v2)
        assert v.shape == out.shape
        assert n.all(v == out)
        # Testing out-of-bounds axis
        with pytest.raises(IndexError):
            stack([a, b], 4)
        with pytest.raises(IndexError):
            stack([a, b], -4)

        # Testing depreciation warning
        with pytest.warns(DeprecationWarning):
            s = stack(a, b)
//  [python]
def test_stack_new_interface(self):
        # Test the new numpy-like interface: stack(tensors, axis=0).

        a = imatrix("a")
        b = imatrix("b")

        # Testing axis parameter
        s3 = stack([a, b], 1)
        f = function([a, b], s3, mode=self.mode)
        v3 = f([[1, 2]], [[3, 4]])
        v4 = np.array([[[1, 2], [3, 4]]])
        assert v3.shape == v4.shape
        assert np.all(v3 == v4)
        # Testing negative axis
        v1 = [[1, 2, 3], [4, 5, 6]]
        v2 = [[7, 8, 9], [10, 11, 12]]
        s = stack([a, b], axis=-1)
        f = function([a, b], s, mode=self.mode)
        v = np.zeros((2, 3, 2))
        v[:, :, 0] = v1
        v[:, :, 1] = v2
        out = f(v1, v2)
        assert v.shape == out.shape
        assert np.all(v == out)
        s = stack([a, b], axis=-2)
        f = function([a, b], s, mode=self.mode)
        v = np.zeros((2, 2, 3))
        v[:, 0, :] = v1
        v[:, 1, :] = v2
        out = f(v1, v2)
        assert v.shape == out.shape
        assert np.all(v == out)
        # Testing out-of-bounds axis
        with pytest.raises(IndexError):
            stack([a, b], 4)
        with pytest.raises(IndexError):
            stack([a, b], -4)

        # Testing depreciation warning
        with pytest.warns(DeprecationWarning):
            s = stack(a, b)
//  [python]
def test_flatten_ndim2_of_3():
    a = TensorType("float64", shape=(None, None, None))()
    c = flatten(a, 2)
    f = inplace_func([a], c)
    a_val = _asarray([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype="float64")
    c_val = _asarray([[0, 1, 2, 3], [4, 5, 6, 7]], dtype="float64")
    assert n.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert n.all(f(a_val) == c_val)

    flatten_2 = partial(flatten, ndim=2)
    utt.verify_grad(flatten_2, [a_val])
//  [python]
def test_flatten_ndim2_of_3():
    a = TensorType("float64", shape=(None, None, None))()
    c = flatten(a, 2)
    f = inplace_func([a], c)
    a_val = _asarray([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype="float64")
    c_val = _asarray([[0, 1, 2, 3], [4, 5, 6, 7]], dtype="float64")
    assert np.all(f(a_val) == c_val)
    f = inplace_func([a], c)
    assert np.all(f(a_val) == c_val)

    flatten_2 = partial(flatten, ndim=2)
    utt.verify_grad(flatten_2, [a_val])
//  [python]
def test_Op_integers(self):
        # Test behaviour of ARange Op on integer inputs
        start, stop, step = iscalars("start", "stop", "step")
        out = ARange(start.type.dtype)(start, stop, step)
        f = function([start, stop, step], out)

        assert n.all(f(0, 5, 1) == n.arange(0, 5, 1))
        assert n.all(f(2, 11, 4) == n.arange(2, 11, 4))
        assert n.all(f(-5, 1, 1) == n.arange(-5, 1, 1))
        assert n.all(f(10, 2, -2) == n.arange(10, 2, -2))
        assert n.all(f(10, 2, 2) == n.arange(10, 2, 2))
        assert n.all(f(0, 0, 1) == n.arange(0, 0, 1))
//  [python]
def test_Op_integers(self):
        # Test behaviour of ARange Op on integer inputs
        start, stop, step = iscalars("start", "stop", "step")
        out = ARange(start.type.dtype)(start, stop, step)
        f = function([start, stop, step], out)

        assert np.all(f(0, 5, 1) == np.arange(0, 5, 1))
        assert np.all(f(2, 11, 4) == np.arange(2, 11, 4))
        assert np.all(f(-5, 1, 1) == np.arange(-5, 1, 1))
        assert np.all(f(10, 2, -2) == np.arange(10, 2, -2))
        assert np.all(f(10, 2, 2) == np.arange(10, 2, 2))
        assert np.all(f(0, 0, 1) == np.arange(0, 0, 1))
//  [python]
def test_tile_grad():
    def grad_tile(x, reps, np_x):
        y = tile(x, reps)
        z = y.sum()
        g = aesara.function([x], grad(z, x))
        grad_res = g(np_x)
        # The gradient should be the product of the tiling dimensions
        # (since the gradients are additive through the tiling operation)
        assert n.all(grad_res == n.prod(reps))

    rng = np.random.default_rng(utt.fetch_seed())

    # test vector
    grad_tile(vector("x"), [3], rng.standard_normal(5).astype(config.floatX))
    # test matrix
    grad_tile(matrix("x"), [3, 4], rng.standard_normal((2, 3)).astype(config.floatX))
    # test tensor3
    grad_tile(
        tensor3("x"), [3, 4, 5], rng.standard_normal((2, 4, 3)).astype(config.floatX)
    )
    # test tensor4
    grad_tile(
        tensor4("x"),
        [3, 4, 5, 6],
        rng.standard_normal((2, 4, 3, 5)).astype(config.floatX),
    )
//  [python]
def test_tile_grad():
    def grad_tile(x, reps, np_x):
        y = tile(x, reps)
        z = y.sum()
        g = aesara.function([x], grad(z, x))
        grad_res = g(np_x)
        # The gradient should be the product of the tiling dimensions
        # (since the gradients are additive through the tiling operation)
        assert np.all(grad_res == np.prod(reps))

    rng = np.random.default_rng(utt.fetch_seed())

    # test vector
    grad_tile(vector("x"), [3], rng.standard_normal(5).astype(config.floatX))
    # test matrix
    grad_tile(matrix("x"), [3, 4], rng.standard_normal((2, 3)).astype(config.floatX))
    # test tensor3
    grad_tile(
        tensor3("x"), [3, 4, 5], rng.standard_normal((2, 4, 3)).astype(config.floatX)
    )
    # test tensor4
    grad_tile(
        tensor4("x"),
        [3, 4, 5, 6],
        rng.standard_normal((2, 4, 3, 5)).astype(config.floatX),
    )
//  [python]
def test_integers(self, cast_policy):
        """Test arange constructor, on integer outputs."""
        with config.change_flags(cast_policy=cast_policy):
            start, stop, step = iscalars("start", "stop", "step")
            out = arange(start, stop, step)
            f = function([start, stop, step], out)

            if cast_policy == "custom":
                assert out.dtype == "int64"
            elif cast_policy == "numpy+floatX":
                numpy_dtype = np.arange(np.array(1, dtype="int32")).dtype
                assert out.dtype == numpy_dtype

            assert n.all(f(0, 5, 1) == np.arange(0, 5, 1))
            assert n.all(f(2, 11, 4) == np.arange(2, 11, 4))
            assert n.all(f(-5, 1, 1) == np.arange(-5, 1, 1))
            assert n.all(f(10, 2, -2) == np.arange(10, 2, -2))
            assert n.all(f(10, 2, 2) == np.arange(10, 2, 2))
            assert n.all(f(0, 0, 1) == np.arange(0, 0, 1))
//  [python]
def test_integers(self, cast_policy):
        """Test arange constructor, on integer outputs."""
        with config.change_flags(cast_policy=cast_policy):
            start, stop, step = iscalars("start", "stop", "step")
            out = arange(start, stop, step)
            f = function([start, stop, step], out)

            if cast_policy == "custom":
                assert out.dtype == "int64"
            elif cast_policy == "numpy+floatX":
                numpy_dtype = np.arange(np.array(1, dtype="int32")).dtype
                assert out.dtype == numpy_dtype

            assert np.all(f(0, 5, 1) == np.arange(0, 5, 1))
            assert np.all(f(2, 11, 4) == np.arange(2, 11, 4))
            assert np.all(f(-5, 1, 1) == np.arange(-5, 1, 1))
            assert np.all(f(10, 2, -2) == np.arange(10, 2, -2))
            assert np.all(f(10, 2, 2) == np.arange(10, 2, 2))
            assert np.all(f(0, 0, 1) == np.arange(0, 0, 1))
//  [python]
def test_float32(self, cast_policy):
        """Test arange constructor, on float32 outputs."""
        with config.change_flags(cast_policy=cast_policy):
            start, stop, step = fscalars("start", "stop", "step")
            out = arange(start, stop, step)
            f = function([start, stop, step], out)

            if config.cast_policy == "custom":
                assert out.dtype == start.type.dtype
            elif config.cast_policy == "numpy+floatX":
                assert out.dtype == config.floatX

            arg_vals = [
                (0, 5, 1),
                (2, 11, 4),
                (-5, 1.1, 1.2),
                (1.3, 2, -2.1),
                (10, 2, 2),
            ]
            for arg_v in arg_vals:
                start_v, stop_v, step_v = arg_v
                start_v_, stop_v_, step_v_ = np.asarray(arg_v, dtype=start.type.dtype)
                f_val = f(start_v_, stop_v_, step_v_)

                if config.cast_policy == "custom":
                    expected_val = np.arange(
                        start_v, stop_v, step_v, dtype=start.type.dtype
                    )
                elif config.cast_policy == "numpy+floatX":
                    expected_val = np.arange(
                        start_v_, stop_v_, step_v_, dtype=out.dtype
                    )

                assert np.all(f_val == expected_val)
//  [python]
def test_float64(self, cast_policy):
        """Test arange constructor, on float64 outputs."""
        with config.change_flags(cast_policy=cast_policy):
            start, stop, step = dscalars("start", "stop", "step")
            out = arange(start, stop, step)
            f = function([start, stop, step], out)

            assert out.dtype == start.type.dtype

            arg_vals = [
                (0, 5, 1),
                (2, 11, 4),
                (-5, 1.1, 1.2),
                (1.3, 2, -2.1),
                (10, 2, 2),
            ]
            for arg_v in arg_vals:
                start_v, stop_v, step_v = arg_v
                start_v_, stop_v_, step_v_ = np.asarray(arg_v, dtype=start.type.dtype)
                f_val = f(start_v_, stop_v_, step_v_)

                if config.cast_policy == "custom":
                    expected_val = np.arange(
                        start_v, stop_v, step_v, dtype=start.type.dtype
                    )
                elif config.cast_policy == "numpy+floatX":
                    expected_val = np.arange(start_v_, stop_v_, step_v_)

                assert np.all(f_val == expected_val)
//  [python]
def test_default_start(self, cast_policy):
        """Test that arange constructor uses the correct default start."""
        with config.change_flags(cast_policy=cast_policy):
            stop = iscalar("stop")
            out = arange(stop)
            f = function([stop], out)

            if config.cast_policy == "custom":
                assert out.dtype == "int64"
            elif config.cast_policy == "numpy+floatX":
                assert out.dtype == np.arange(np.int32(1)).dtype

            assert n.all(f(8) == np.arange(8))
            assert n.all(f(-2) == np.arange(-2))

            fstop = fscalar("stop")
            fout = arange(fstop)
            ff = function([fstop], fout)

            if config.cast_policy == "custom":
                assert fout.dtype == fstop.type.dtype
            elif config.cast_policy == "numpy+floatX":
                if config.floatX == "float32":
                    assert fout.dtype == "float32"
                else:
                    assert fout.dtype == np.arange(np.float32(1)).dtype

            fstop_values = [0.2, -0.7, 8.5]
            for fstop_v in fstop_values:
                fstop_v32 = np.float32(fstop_v)
                assert n.all(ff(fstop_v32) == np.arange(fstop_v))
//  [python]
def test_default_start(self, cast_policy):
        """Test that arange constructor uses the correct default start."""
        with config.change_flags(cast_policy=cast_policy):
            stop = iscalar("stop")
            out = arange(stop)
            f = function([stop], out)

            if config.cast_policy == "custom":
                assert out.dtype == "int64"
            elif config.cast_policy == "numpy+floatX":
                assert out.dtype == np.arange(np.int32(1)).dtype

            assert np.all(f(8) == np.arange(8))
            assert np.all(f(-2) == np.arange(-2))

            fstop = fscalar("stop")
            fout = arange(fstop)
            ff = function([fstop], fout)

            if config.cast_policy == "custom":
                assert fout.dtype == fstop.type.dtype
            elif config.cast_policy == "numpy+floatX":
                if config.floatX == "float32":
                    assert fout.dtype == "float32"
                else:
                    assert fout.dtype == np.arange(np.float32(1)).dtype

            fstop_values = [0.2, -0.7, 8.5]
            for fstop_v in fstop_values:
                fstop_v32 = np.float32(fstop_v)
                assert np.all(ff(fstop_v32) == np.arange(fstop_v))
//  [python]
def test_default_step(self, cast_policy):
        """Test that arange constructor uses the correct default step."""
        with config.change_flags(cast_policy=cast_policy):
            start, stop = iscalars("start", "stop")
            out = arange(start, stop)
            f = function([start, stop], out)

            if config.cast_policy == "custom":
                assert out.dtype == "int64"
            elif config.cast_policy == "numpy+floatX":
                assert out.dtype == np.arange(np.int32(0), np.int32(1)).dtype

            assert n.all(f(0, 5) == np.arange(0, 5))
            assert n.all(f(-5, 1) == np.arange(-5, 1))
            assert n.all(f(0, 0) == np.arange(0, 0))

            dstart, dstop = dscalars("start", "stop")
            dout = arange(dstart, dstop)
            df = function([dstart, dstop], dout)

            assert dout.dtype == dstart.type.dtype
            # print df(0.2, 5.3)
            # print np.arange(0.2, 5.3)
            assert n.all(df(0.2, 5.3) == np.arange(0.2, 5.3))
            assert n.all(df(0.8, 5.3) == np.arange(0.8, 5.3))
            assert n.all(df(-0.7, 5.3) == np.arange(-0.7, 5.3))
//  [python]
def test_default_step(self, cast_policy):
        """Test that arange constructor uses the correct default step."""
        with config.change_flags(cast_policy=cast_policy):
            start, stop = iscalars("start", "stop")
            out = arange(start, stop)
            f = function([start, stop], out)

            if config.cast_policy == "custom":
                assert out.dtype == "int64"
            elif config.cast_policy == "numpy+floatX":
                assert out.dtype == np.arange(np.int32(0), np.int32(1)).dtype

            assert np.all(f(0, 5) == np.arange(0, 5))
            assert np.all(f(-5, 1) == np.arange(-5, 1))
            assert np.all(f(0, 0) == np.arange(0, 0))

            dstart, dstop = dscalars("start", "stop")
            dout = arange(dstart, dstop)
            df = function([dstart, dstop], dout)

            assert dout.dtype == dstart.type.dtype
            # print df(0.2, 5.3)
            # print np.arange(0.2, 5.3)
            assert np.all(df(0.2, 5.3) == np.arange(0.2, 5.3))
            assert np.all(df(0.8, 5.3) == np.arange(0.8, 5.3))
            assert np.all(df(-0.7, 5.3) == np.arange(-0.7, 5.3))
//  [python]
def test_dim2(self):
        # Test the inversion of several permutations at a time
        # Each row of p is a different permutation to inverse
        p = imatrix()
        inv = inverse_permutation(p)
        f_inverse = function([p], inv)

        rng = np.random.default_rng(utt.fetch_seed())
        # Generate 10 random permutations
        p_val = np.asarray([rng.permutation(10) for i in range(7)], dtype="int32")
        inv_val = f_inverse(p_val)

        # Check that the inverse of the inverse is the original permutation list
        assert n.all(f_inverse(inv_val) == p_val)
        # Check that, for each permutation,
        # permutation(inverse) == inverse(permutation) = identity
        for p_row, i_row in zip(p_val, inv_val):
            assert n.all(p_row[i_row] == np.arange(10))
            assert n.all(i_row[p_row] == np.arange(10))
//  [python]
def test_dim2(self):
        # Test the inversion of several permutations at a time
        # Each row of p is a different permutation to inverse
        p = imatrix()
        inv = inverse_permutation(p)
        f_inverse = function([p], inv)

        rng = np.random.default_rng(utt.fetch_seed())
        # Generate 10 random permutations
        p_val = np.asarray([rng.permutation(10) for i in range(7)], dtype="int32")
        inv_val = f_inverse(p_val)

        # Check that the inverse of the inverse is the original permutation list
        assert np.all(f_inverse(inv_val) == p_val)
        # Check that, for each permutation,
        # permutation(inverse) == inverse(permutation) = identity
        for p_row, i_row in zip(p_val, inv_val):
            assert np.all(p_row[i_row] == np.arange(10))
            assert np.all(i_row[p_row] == np.arange(10))
//  [python]
def test_dim1(self):
        # Test the inversion of one permutation (int vector)
        p = ivector()
        inv = inverse_permutation(p)
        assert inv.dtype == p.dtype
        f_inverse = function([p], inv)

        # Generate a random permutation
        rng = np.random.default_rng(utt.fetch_seed())
        p_val = rng.permutation(10).astype("int32")
        inv_val = f_inverse(p_val)

        # Check that the inverse of the inverse is the original permutation
        assert n.all(f_inverse(inv_val) == p_val)
        # Check that permutation(inverse) == inverse(permutation) = identity
        assert n.all(p_val[inv_val] == n.arange(10))
        assert n.all(inv_val[p_val] == n.arange(10))

        # Test passing a list
        p = [2, 4, 3, 0, 1]
        inv = at.inverse_permutation(p)
        f = aesara.function([], inv)
        assert n.array_equal(f(), n.array([3, 4, 0, 2, 1]))
//  [python]
def test_dim1(self):
        # Test the inversion of one permutation (int vector)
        p = ivector()
        inv = inverse_permutation(p)
        assert inv.dtype == p.dtype
        f_inverse = function([p], inv)

        # Generate a random permutation
        rng = np.random.default_rng(utt.fetch_seed())
        p_val = rng.permutation(10).astype("int32")
        inv_val = f_inverse(p_val)

        # Check that the inverse of the inverse is the original permutation
        assert np.all(f_inverse(inv_val) == p_val)
        # Check that permutation(inverse) == inverse(permutation) = identity
        assert np.all(p_val[inv_val] == np.arange(10))
        assert np.all(inv_val[p_val] == np.arange(10))

        # Test passing a list
        p = [2, 4, 3, 0, 1]
        inv = at.inverse_permutation(p)
        f = aesara.function([], inv)
        assert np.array_equal(f(), np.array([3, 4, 0, 2, 1]))
//  [python]
def test_1_1(self):
        # Test PermuteRowElements(vector, vector)
        input = dvector()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(5,))
        p_val = rng.permutation(5).astype("int32")
        out_val = permute(input_val, p_val)

        # Should be equivalent to advanced indexing
        out_bis = input_val[p_val]
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_1_1(self):
        # Test PermuteRowElements(vector, vector)
        input = dvector()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(5,))
        p_val = rng.permutation(5).astype("int32")
        out_val = permute(input_val, p_val)

        # Should be equivalent to advanced indexing
        out_bis = input_val[p_val]
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_1(self):
        # Test broadcasting in PermuteRowElements(matrix, vector)
        input = matrix()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = rng.permutation(5).astype("int32")
        out_val = permute(input_val, p_val)

        # The same permutation should be applied to every row of the input matrix.
        out_bis = np.asarray([r[p_val] for r in input_val])
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_1(self):
        # Test broadcasting in PermuteRowElements(matrix, vector)
        input = matrix()
        p = ivector()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = rng.permutation(5).astype("int32")
        out_val = permute(input_val, p_val)

        # The same permutation should be applied to every row of the input matrix.
        out_bis = np.asarray([r[p_val] for r in input_val])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_1_2(self):
        # Test PermuteRowElements(vector, matrix)
        # Different permutations will be applied to the same input vector
        input = vector()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(5,)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)], dtype="int32")
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the input vector
        out_bis = np.asarray([input_val[p_row] for p_row in p_val])
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_1_2(self):
        # Test PermuteRowElements(vector, matrix)
        # Different permutations will be applied to the same input vector
        input = vector()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(5,)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)], dtype="int32")
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the input vector
        out_bis = np.asarray([input_val[p_row] for p_row in p_val])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_2_2(self):
        # Test PermuteRowElements(matrix, matrix)
        input = matrix()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(3, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)], dtype="int32")
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to the corresponding
        # row of input
        out_bis = np.asarray([i_row[p_row] for i_row, p_row in zip(input_val, p_val)])
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_fit_int64(self):
        bitwidth = PYTHON_INT_BITWIDTH
        for exponent in range(bitwidth):
            val = 2**exponent - 1
            scalar_ct = constant(val)

            assert scalar_ct.dtype in int_dtypes, (
                exponent,
                val,
                scalar_ct.dtype,
            )
            assert scalar_ct.value == val

            vector_ct = constant([val, val])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if bitwidth == 32:
                assert vector_ct.dtype == "int32"
            else:
                assert vector_ct.dtype == "int64"
            assert n.all(vector_ct.value == val)

            matrix_ct = constant([[val, val]])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if bitwidth == 32:
                assert matrix_ct.dtype == "int32"
            else:
                assert matrix_ct.dtype == "int64"
            assert n.all(matrix_ct.value == val)
//  [python]
def test_fit_int64(self):
        bitwidth = PYTHON_INT_BITWIDTH
        for exponent in range(bitwidth):
            val = 2**exponent - 1
            scalar_ct = constant(val)

            assert scalar_ct.dtype in int_dtypes, (
                exponent,
                val,
                scalar_ct.dtype,
            )
            assert scalar_ct.value == val

            vector_ct = constant([val, val])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if bitwidth == 32:
                assert vector_ct.dtype == "int32"
            else:
                assert vector_ct.dtype == "int64"
            assert np.all(vector_ct.value == val)

            matrix_ct = constant([[val, val]])
            # On Python 2, np.array() on a "long" returns int64,
            # but on Python 3, all integers are long, and np.asarray
            # will not force the upcasting, and return the native int width.
            if bitwidth == 32:
                assert matrix_ct.dtype == "int32"
            else:
                assert matrix_ct.dtype == "int64"
            assert np.all(matrix_ct.value == val)
//  [python]
def test_3b_2(self):
        # Test permute_row_elements on a more complex broadcasting pattern:
        # input.type.shape = (None, 1, None),
        # p.type.shape = (None, None).

        input = TensorType("floatX", shape=(None, 1, None))()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(4, 1, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)], dtype="int32")
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to each row
        # of the input tensor
        out_bis = np.asarray(
            [[in_mat[0, p_row] for p_row in p_val] for in_mat in input_val]
        )
        assert n.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def test_3b_2(self):
        # Test permute_row_elements on a more complex broadcasting pattern:
        # input.type.shape = (None, 1, None),
        # p.type.shape = (None, None).

        input = TensorType("floatX", shape=(None, 1, None))()
        p = imatrix()
        out = permute_row_elements(input, p)
        permute = function([input, p], out)

        rng = np.random.default_rng(utt.fetch_seed())
        input_val = rng.uniform(size=(4, 1, 5)).astype(config.floatX)
        p_val = np.asarray([rng.permutation(5) for i in range(3)], dtype="int32")
        out_val = permute(input_val, p_val)

        # Each row of p contains a permutation to apply to each row
        # of the input tensor
        out_bis = np.asarray(
            [[in_mat[0, p_row] for p_row in p_val] for in_mat in input_val]
        )
        assert np.all(out_val == out_bis)

        # Verify gradient
        def permute_fixed(s_input):
            # Auxiliary op defined to get rid of gradient wrt p_val
            return permute_row_elements(s_input, p_val)

        utt.verify_grad(permute_fixed, [input_val])
//  [python]
def isna(self) -> np.array:
        """
        See docstring in :class:`ExtensionArray` class in ``pandas/core/arrays/base.py``
        for information about this method.
        """
        if self._tensor.dtype.type is np.object_:
            # Avoid comparing with __eq__ because the elements of the tensor may do
            # something funny with that operation.
            result_list = [
                self._tensor[i] is None for i in range(len(self))
            ]
            return np.array(result_list, dtype=bool)
        elif self._tensor.dtype.type is np.str_:
            return np.all(self._tensor == "", axis=-1)
        else:
            return np.isnan(self._tensor).all(axis=-1)
//  [python]
def isna(self) -> np.array:
        """
        See docstring in :class:`ExtensionArray` class in ``pandas/core/arrays/base.py``
        for information about this method.
        """
        if self._tensor.dtype.type is np.object_:
            # Avoid comparing with __eq__ because the elements of the tensor may do
            # something funny with that operation.
            result_list = [
                self._tensor[i] is None for i in range(len(self))
            ]
            return np.array(result_list, dtype=bool)
        elif self._tensor.dtype.type is np.str_:
            return np.all(self._tensor == "", axis=-1)
        else:
            return np.all(np.isnan(self._tensor), axis=-1)
//  [python]
def test_same_seed(self, measurements, max_workers):
        """Test that different devices given the same random seed will produce
        the same results"""
        qs = qml.tape.QuantumScript([qml.Hadamard(0)], measurements, shots=1000)

        dev1 = DefaultQubit(seed=123, max_workers=max_workers)
        result1 = dev1.execute(qs)

        dev2 = DefaultQubit(seed=123, max_workers=max_workers)
        result2 = dev2.execute(qs)

        if len(measurements) == 1:
            result1, result2 = [result1], [result2]

        assert all(n.all(res1 == res2) for res1, res2 in zip(result1, result2))
//  [python]
def test_same_seed(self, measurements, max_workers):
        """Test that different devices given the same random seed will produce
        the same results"""
        qs = qml.tape.QuantumScript([qml.Hadamard(0)], measurements, shots=1000)

        dev1 = DefaultQubit(seed=123, max_workers=max_workers)
        result1 = dev1.execute(qs)

        dev2 = DefaultQubit(seed=123, max_workers=max_workers)
        result2 = dev2.execute(qs)

        if len(measurements) == 1:
            result1, result2 = [result1], [result2]

        assert all(np.all(res1 == res2) for res1, res2 in zip(result1, result2))