jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_permutation(self):
        """Test that raw_random.permutation generates the same
        results as numpy."""
        rng_R = random_state_type()
        post_r, out = permutation(rng_R, size=(9,), n=6)
        f = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r, mutable=True)],
                [out], accept_inplace=True)

        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        # Check over two calls to see if the random state is correctly updated.
        # numpy_rng.permutation outputs one vector at a time,
        # so we call it iteratively to generate all the samples.
        val0 = f()
        val1 = f()
        numpy_val0 = numpy.asarray([numpy_rng.permutation(6)
                                    for i in range(9)])
        numpy_val1 = numpy.asarray([numpy_rng.permutation(6)
                                    for i in range(9)])
        self.assertTrue(numpy.all(val0 == numpy_val0))
        self.assertTrue(numpy.all(val1 == numpy_val1))

        # Test that we can generate a list: have size=None or ().
        for ndim in [1, None]:
            post_r, out = permutation(rng_R, n=10, size=None, ndim=ndim)
            inp = compile.In(rng_R,
                             value=numpy.random.RandomState(utt.fetch_seed()),
                             update=post_r, mutable=True)
            f = theano.function([inp], out)
            o = f()
            assert o.shape == (10,)
            assert (numpy.sort(o) == numpy.arange(10)).all()
        # Wrong number of dimensions asked
        self.assertRaises(TypeError, permutation, rng_R, size=None, ndim=2)
//  [python]
def test_broadcast_arguments(self):
        rng_R = random_state_type()
        low = tensor.dvector()
        high = tensor.dcol()
        post_r, out = uniform(rng_R, low=low, high=high)
        assert out.ndim == 2
        f = compile.function([rng_R, low, high], [post_r, out],
                             accept_inplace=True)

        rng_state0 = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        post0, val0 = f(rng_state0, [-5, .5, 0, 1], [[1.]])
        post1, val1 = f(post0, [.9], [[1.], [1.1], [1.5]])
        post2, val2 = f(post1, [-5, .5, 0, 1], [[1.], [1.1], [1.5]])

        numpy_val0 = numpy_rng.uniform(low=[-5, .5, 0, 1], high=[1.])
        numpy_val1 = numpy_rng.uniform(low=[.9], high=[[1.], [1.1], [1.5]])
        numpy_val2 = numpy_rng.uniform(low=[-5, .5, 0, 1],
                                       high=[[1.], [1.1], [1.5]])

        assert n.all(val0 == numpy_val0), (val0, numpy_val0)
        assert n.all(val1 == numpy_val1)
        assert n.all(val2 == numpy_val2)
//  [python]
def test_broadcast_arguments(self):
        rng_R = random_state_type()
        low = tensor.dvector()
        high = tensor.dcol()
        post_r, out = uniform(rng_R, low=low, high=high)
        assert out.ndim == 2
        f = compile.function([rng_R, low, high], [post_r, out],
                             accept_inplace=True)

        rng_state0 = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        post0, val0 = f(rng_state0, [-5, .5, 0, 1], [[1.]])
        post1, val1 = f(post0, [.9], [[1.], [1.1], [1.5]])
        post2, val2 = f(post1, [-5, .5, 0, 1], [[1.], [1.1], [1.5]])

        numpy_val0 = numpy_rng.uniform(low=[-5, .5, 0, 1], high=[1.])
        numpy_val1 = numpy_rng.uniform(low=[.9], high=[[1.], [1.1], [1.5]])
        numpy_val2 = numpy_rng.uniform(low=[-5, .5, 0, 1],
                                       high=[[1.], [1.1], [1.5]])

        assert numpy.all(val0 == numpy_val0), (val0, numpy_val0)
        assert numpy.all(val1 == numpy_val1)
        assert numpy.all(val2 == numpy_val2)
//  [python]
def test_permutation_helper(self):
        """Test that raw_random.permutation_helper generates the same
        results as numpy,
        and that the 'ndim_added' keyword behaves correctly."""
        # permutation_helper needs "ndim_added=1", because its output
        # is one dimension more than its "shape" argument (and there's
        # no way to determine that automatically).
        # Check the working case, over two calls to see if the random
        # state is correctly updated.
        rf = RandomFunction(permutation_helper, tensor.imatrix, 8,
                            ndim_added=1)
        rng_R = random_state_type()
        post_r, out = rf(rng_R, (7,), 8)

        f = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r, mutable=True)],
                [out], accept_inplace=True)

        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        val0 = f()
        val1 = f()
        # numpy_rng.permutation outputs one vector at a time,
        # so we call it iteratively to generate all the samples.
        numpy_val0 = numpy.asarray([numpy_rng.permutation(8)
                                    for i in range(7)])
        numpy_val1 = numpy.asarray([numpy_rng.permutation(8)
                                    for i in range(7)])
        self.assertTrue(n.all(val0 == numpy_val0))
        self.assertTrue(n.all(val1 == numpy_val1))

        # This call lacks "ndim_added=1", so ndim_added defaults to 0.
        # A ValueError should be raised.
        rf0 = RandomFunction(permutation_helper, tensor.imatrix, 8)
        post_r0, out0 = rf0(rng_R, (7,), 8)
        f0 = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r0, mutable=True)],
                [out0], accept_inplace=True)
        self.assertRaises(ValueError, f0)

        # Here, ndim_added is 2 instead of 1. A ValueError should be raised.
        rf2 = RandomFunction(permutation_helper, tensor.imatrix, 8,
                             ndim_added=2)
        post_r2, out2 = rf2(rng_R, (7,), 8)
        f2 = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r2, mutable=True)],
                [out2], accept_inplace=True)
        self.assertRaises(ValueError, f2)
//  [python]
def test_permutation_helper(self):
        """Test that raw_random.permutation_helper generates the same
        results as numpy,
        and that the 'ndim_added' keyword behaves correctly."""
        # permutation_helper needs "ndim_added=1", because its output
        # is one dimension more than its "shape" argument (and there's
        # no way to determine that automatically).
        # Check the working case, over two calls to see if the random
        # state is correctly updated.
        rf = RandomFunction(permutation_helper, tensor.imatrix, 8,
                            ndim_added=1)
        rng_R = random_state_type()
        post_r, out = rf(rng_R, (7,), 8)

        f = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r, mutable=True)],
                [out], accept_inplace=True)

        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        val0 = f()
        val1 = f()
        # numpy_rng.permutation outputs one vector at a time,
        # so we call it iteratively to generate all the samples.
        numpy_val0 = numpy.asarray([numpy_rng.permutation(8)
                                    for i in range(7)])
        numpy_val1 = numpy.asarray([numpy_rng.permutation(8)
                                    for i in range(7)])
        self.assertTrue(numpy.all(val0 == numpy_val0))
        self.assertTrue(numpy.all(val1 == numpy_val1))

        # This call lacks "ndim_added=1", so ndim_added defaults to 0.
        # A ValueError should be raised.
        rf0 = RandomFunction(permutation_helper, tensor.imatrix, 8)
        post_r0, out0 = rf0(rng_R, (7,), 8)
        f0 = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r0, mutable=True)],
                [out0], accept_inplace=True)
        self.assertRaises(ValueError, f0)

        # Here, ndim_added is 2 instead of 1. A ValueError should be raised.
        rf2 = RandomFunction(permutation_helper, tensor.imatrix, 8,
                             ndim_added=2)
        post_r2, out2 = rf2(rng_R, (7,), 8)
        f2 = compile.function(
                [compile.In(rng_R,
                    value=numpy.random.RandomState(utt.fetch_seed()),
                    update=post_r2, mutable=True)],
                [out2], accept_inplace=True)
        self.assertRaises(ValueError, f2)
//  [python]
def test_uniform_vector(self):
        rng_R = random_state_type()
        low = tensor.vector()
        high = tensor.vector()
        post_r, out = uniform(rng_R, low=low, high=high)
        assert out.ndim == 1
        f = compile.function([rng_R, low, high], [post_r, out],
                             accept_inplace=True)

        def as_floatX(thing):
            return numpy.asarray(thing, dtype=theano.config.floatX)
        low_val = as_floatX([.1, .2, .3])
        high_val = as_floatX([1.1, 2.2, 3.3])
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, low_val, high_val)
        numpy_val0 = as_floatX(numpy_rng.uniform(low=low_val, high=high_val))
        assert n.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, low_val[:-1], high_val[:-1])
        numpy_val1 = as_floatX(numpy_rng.uniform(low=low_val[:-1],
                                                 high=high_val[:-1]))
        assert n.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, low, high],
                uniform(rng_R, low=low, high=high, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, low_val, high_val)
        numpy_val2 = as_floatX(numpy_rng.uniform(low=low_val, high=high_val,
                                                 size=(3,)))
        assert n.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, low_val[:-1], high_val[:-1])
//  [python]
def test_uniform_vector(self):
        rng_R = random_state_type()
        low = tensor.vector()
        high = tensor.vector()
        post_r, out = uniform(rng_R, low=low, high=high)
        assert out.ndim == 1
        f = compile.function([rng_R, low, high], [post_r, out],
                             accept_inplace=True)

        def as_floatX(thing):
            return numpy.asarray(thing, dtype=theano.config.floatX)
        low_val = as_floatX([.1, .2, .3])
        high_val = as_floatX([1.1, 2.2, 3.3])
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, low_val, high_val)
        numpy_val0 = as_floatX(numpy_rng.uniform(low=low_val, high=high_val))
        assert numpy.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, low_val[:-1], high_val[:-1])
        numpy_val1 = as_floatX(numpy_rng.uniform(low=low_val[:-1],
                                                 high=high_val[:-1]))
        assert numpy.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, low, high],
                uniform(rng_R, low=low, high=high, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, low_val, high_val)
        numpy_val2 = as_floatX(numpy_rng.uniform(low=low_val, high=high_val,
                                                 size=(3,)))
        assert numpy.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, low_val[:-1], high_val[:-1])
//  [python]
def test_binomial_vector(self):
        rng_R = random_state_type()
        n = tensor.lvector()
        prob = tensor.vector()
        post_r, out = binomial(rng_R, n=n, p=prob)
        assert out.ndim == 1
        f = compile.function([rng_R, n, prob], [post_r, out],
                             accept_inplace=True)

        n_val = [1, 2, 3]
        prob_val = numpy.asarray([.1, .2, .3], dtype=config.floatX)
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, n_val, prob_val)
        numpy_val0 = numpy_rng.binomial(n=n_val, p=prob_val)
        assert n.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, n_val[:-1], prob_val[:-1])
        numpy_val1 = numpy_rng.binomial(n=n_val[:-1], p=prob_val[:-1])
        assert n.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, n, prob],
                binomial(rng_R, n=n, p=prob, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, n_val, prob_val)
        numpy_val2 = numpy_rng.binomial(n=n_val, p=prob_val, size=(3,))
        assert n.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, n_val[:-1], prob_val[:-1])
//  [python]
def test_binomial_vector(self):
        rng_R = random_state_type()
        n = tensor.lvector()
        prob = tensor.vector()
        post_r, out = binomial(rng_R, n=n, p=prob)
        assert out.ndim == 1
        f = compile.function([rng_R, n, prob], [post_r, out],
                             accept_inplace=True)

        n_val = [1, 2, 3]
        prob_val = numpy.asarray([.1, .2, .3], dtype=config.floatX)
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, n_val, prob_val)
        numpy_val0 = numpy_rng.binomial(n=n_val, p=prob_val)
        assert numpy.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, n_val[:-1], prob_val[:-1])
        numpy_val1 = numpy_rng.binomial(n=n_val[:-1], p=prob_val[:-1])
        assert numpy.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, n, prob],
                binomial(rng_R, n=n, p=prob, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, n_val, prob_val)
        numpy_val2 = numpy_rng.binomial(n=n_val, p=prob_val, size=(3,))
        assert numpy.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, n_val[:-1], prob_val[:-1])
//  [python]
def test_dtype(self):
        rng_R = random_state_type()
        low = tensor.lscalar()
        high = tensor.lscalar()
        post_r, out = random_integers(rng_R, low=low, high=high, size=(20, ),
                                      dtype='int8')
        assert out.dtype == 'int8'
        f = compile.function([rng_R, low, high], [post_r, out])

        rng = numpy.random.RandomState(utt.fetch_seed())
        rng0, val0 = f(rng, 0, 9)
        assert val0.dtype == 'int8'

        rng1, val1 = f(rng0, 255, 257)
        assert val1.dtype == 'int8'
        assert n.all(abs(val1) <= 1)
//  [python]
def test_dtype(self):
        rng_R = random_state_type()
        low = tensor.lscalar()
        high = tensor.lscalar()
        post_r, out = random_integers(rng_R, low=low, high=high, size=(20, ),
                                      dtype='int8')
        assert out.dtype == 'int8'
        f = compile.function([rng_R, low, high], [post_r, out])

        rng = numpy.random.RandomState(utt.fetch_seed())
        rng0, val0 = f(rng, 0, 9)
        assert val0.dtype == 'int8'

        rng1, val1 = f(rng0, 255, 257)
        assert val1.dtype == 'int8'
        assert numpy.all(abs(val1) <= 1)
//  [python]
def test_vector_arguments(self):
        rng_R = random_state_type()
        low = tensor.vector()
        post_r, out = uniform(rng_R, low=low, high=1)
        assert out.ndim == 1
        f = compile.function([rng_R, low], [post_r, out], accept_inplace=True)

        def as_floatX(thing):
            return numpy.asarray(thing, dtype=theano.config.floatX)

        rng_state0 = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        post0, val0 = f(rng_state0, [-5, .5, 0, 1])
        post1, val1 = f(post0, as_floatX([.9]))
        numpy_val0 = as_floatX(numpy_rng.uniform(low=[-5, .5, 0, 1], high=1))
        numpy_val1 = as_floatX(numpy_rng.uniform(low=as_floatX([.9]), high=1))

        assert n.all(val0 == numpy_val0)
        assert n.all(val1 == numpy_val1)

        high = tensor.vector()
        post_rb, outb = uniform(rng_R, low=low, high=high)
        assert outb.ndim == 1
        fb = compile.function([rng_R, low, high], [post_rb, outb],
                              accept_inplace=True)

        post0b, val0b = fb(post1, [-4., -2], [-1, 0])
        post1b, val1b = fb(post0b, [-4.], [-1])
        numpy_val0b = as_floatX(numpy_rng.uniform(low=[-4., -2], high=[-1, 0]))
        numpy_val1b = as_floatX(numpy_rng.uniform(low=[-4.], high=[-1]))
        assert n.all(val0b == numpy_val0b)
        assert n.all(val1b == numpy_val1b)
        self.assertRaises(ValueError, fb, post1b, [-4., -2], [-1, 0, 1])
        # TODO: do we want that?
        #self.assertRaises(ValueError, fb, post1b, [-4., -2], [-1])

        size = tensor.lvector()
        post_rc, outc = uniform(rng_R, low=low, high=high, size=size, ndim=1)
        fc = compile.function([rng_R, low, high, size], [post_rc, outc],
                              accept_inplace=True)
        post0c, val0c = fc(post1b, [-4., -2], [-1, 0], [2])
        post1c, val1c = fc(post0c, [-4.], [-1], [1])
        numpy_val0c = as_floatX(numpy_rng.uniform(low=[-4., -2], high=[-1, 0]))
        numpy_val1c = as_floatX(numpy_rng.uniform(low=[-4.], high=[-1]))
        assert n.all(val0c == numpy_val0c)
        assert n.all(val1c == numpy_val1c)
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1, 0], [1])
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1, 0], [1, 2])
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1, 0], [2, 1])
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1], [1])
//  [python]
def test_vector_arguments(self):
        rng_R = random_state_type()
        low = tensor.vector()
        post_r, out = uniform(rng_R, low=low, high=1)
        assert out.ndim == 1
        f = compile.function([rng_R, low], [post_r, out], accept_inplace=True)

        def as_floatX(thing):
            return numpy.asarray(thing, dtype=theano.config.floatX)

        rng_state0 = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())
        post0, val0 = f(rng_state0, [-5, .5, 0, 1])
        post1, val1 = f(post0, as_floatX([.9]))
        numpy_val0 = as_floatX(numpy_rng.uniform(low=[-5, .5, 0, 1], high=1))
        numpy_val1 = as_floatX(numpy_rng.uniform(low=as_floatX([.9]), high=1))

        assert numpy.all(val0 == numpy_val0)
        assert numpy.all(val1 == numpy_val1)

        high = tensor.vector()
        post_rb, outb = uniform(rng_R, low=low, high=high)
        assert outb.ndim == 1
        fb = compile.function([rng_R, low, high], [post_rb, outb],
                              accept_inplace=True)

        post0b, val0b = fb(post1, [-4., -2], [-1, 0])
        post1b, val1b = fb(post0b, [-4.], [-1])
        numpy_val0b = as_floatX(numpy_rng.uniform(low=[-4., -2], high=[-1, 0]))
        numpy_val1b = as_floatX(numpy_rng.uniform(low=[-4.], high=[-1]))
        assert numpy.all(val0b == numpy_val0b)
        assert numpy.all(val1b == numpy_val1b)
        self.assertRaises(ValueError, fb, post1b, [-4., -2], [-1, 0, 1])
        # TODO: do we want that?
        #self.assertRaises(ValueError, fb, post1b, [-4., -2], [-1])

        size = tensor.lvector()
        post_rc, outc = uniform(rng_R, low=low, high=high, size=size, ndim=1)
        fc = compile.function([rng_R, low, high, size], [post_rc, outc],
                              accept_inplace=True)
        post0c, val0c = fc(post1b, [-4., -2], [-1, 0], [2])
        post1c, val1c = fc(post0c, [-4.], [-1], [1])
        numpy_val0c = as_floatX(numpy_rng.uniform(low=[-4., -2], high=[-1, 0]))
        numpy_val1c = as_floatX(numpy_rng.uniform(low=[-4.], high=[-1]))
        assert numpy.all(val0c == numpy_val0c)
        assert numpy.all(val1c == numpy_val1c)
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1, 0], [1])
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1, 0], [1, 2])
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1, 0], [2, 1])
        self.assertRaises(ValueError, fc, post1c, [-4., -2], [-1], [1])
//  [python]
def test_random_integers_vector(self):
        rng_R = random_state_type()
        low = tensor.lvector()
        high = tensor.lvector()
        post_r, out = random_integers(rng_R, low=low, high=high)
        assert out.ndim == 1
        f = compile.function([rng_R, low, high], [post_r, out],
                             accept_inplace=True)

        low_val = [100, 200, 300]
        high_val = [110, 220, 330]
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, low_val, high_val)
        numpy_val0 = numpy.asarray([numpy_rng.random_integers(low=lv, high=hv)
            for lv, hv in zip(low_val, high_val)])
        assert n.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, low_val[:-1], high_val[:-1])
        numpy_val1 = numpy.asarray([numpy_rng.random_integers(low=lv, high=hv)
            for lv, hv in zip(low_val[:-1], high_val[:-1])])
        assert n.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, low, high],
                random_integers(rng_R, low=low, high=high, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, low_val, high_val)
        numpy_val2 = numpy.asarray([numpy_rng.random_integers(low=lv, high=hv)
            for lv, hv in zip(low_val, high_val)])
        assert n.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, low_val[:-1], high_val[:-1])
//  [python]
def test_random_integers_vector(self):
        rng_R = random_state_type()
        low = tensor.lvector()
        high = tensor.lvector()
        post_r, out = random_integers(rng_R, low=low, high=high)
        assert out.ndim == 1
        f = compile.function([rng_R, low, high], [post_r, out],
                             accept_inplace=True)

        low_val = [100, 200, 300]
        high_val = [110, 220, 330]
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, low_val, high_val)
        numpy_val0 = numpy.asarray([numpy_rng.random_integers(low=lv, high=hv)
            for lv, hv in zip(low_val, high_val)])
        assert numpy.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, low_val[:-1], high_val[:-1])
        numpy_val1 = numpy.asarray([numpy_rng.random_integers(low=lv, high=hv)
            for lv, hv in zip(low_val[:-1], high_val[:-1])])
        assert numpy.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, low, high],
                random_integers(rng_R, low=low, high=high, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, low_val, high_val)
        numpy_val2 = numpy.asarray([numpy_rng.random_integers(low=lv, high=hv)
            for lv, hv in zip(low_val, high_val)])
        assert numpy.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, low_val[:-1], high_val[:-1])
//  [python]
def test_normal_vector(self):
        rng_R = random_state_type()
        avg = tensor.vector()
        std = tensor.vector()
        post_r, out = normal(rng_R, avg=avg, std=std)
        assert out.ndim == 1
        f = compile.function([rng_R, avg, std], [post_r, out],
                             accept_inplace=True)

        def as_floatX(thing):
            return numpy.asarray(thing, dtype=theano.config.floatX)

        avg_val = [1, 2, 3]
        std_val = as_floatX([.1, .2, .3])
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, avg_val, std_val)
        numpy_val0 = as_floatX(numpy_rng.normal(loc=as_floatX(avg_val),
            scale=as_floatX(std_val)))
        assert n.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, avg_val[:-1], std_val[:-1])
        numpy_val1 = numpy.asarray(numpy_rng.normal(loc=avg_val[:-1],
                                                    scale=std_val[:-1]),
                dtype=theano.config.floatX)
        assert n.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, avg, std],
                normal(rng_R, avg=avg, std=std, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, avg_val, std_val)
        numpy_val2 = numpy.asarray(numpy_rng.normal(loc=avg_val, scale=std_val,
                                                    size=(3,)),
                dtype=theano.config.floatX)
        assert n.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, avg_val[:-1], std_val[:-1])
//  [python]
def test_normal_vector(self):
        rng_R = random_state_type()
        avg = tensor.vector()
        std = tensor.vector()
        post_r, out = normal(rng_R, avg=avg, std=std)
        assert out.ndim == 1
        f = compile.function([rng_R, avg, std], [post_r, out],
                             accept_inplace=True)

        def as_floatX(thing):
            return numpy.asarray(thing, dtype=theano.config.floatX)

        avg_val = [1, 2, 3]
        std_val = as_floatX([.1, .2, .3])
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, avg_val, std_val)
        numpy_val0 = as_floatX(numpy_rng.normal(loc=as_floatX(avg_val),
            scale=as_floatX(std_val)))
        assert numpy.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, avg_val[:-1], std_val[:-1])
        numpy_val1 = numpy.asarray(numpy_rng.normal(loc=avg_val[:-1],
                                                    scale=std_val[:-1]),
                dtype=theano.config.floatX)
        assert numpy.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, avg, std],
                normal(rng_R, avg=avg, std=std, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, avg_val, std_val)
        numpy_val2 = numpy.asarray(numpy_rng.normal(loc=avg_val, scale=std_val,
                                                    size=(3,)),
                dtype=theano.config.floatX)
        assert numpy.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, avg_val[:-1], std_val[:-1])
//  [python]
def test_any_all_object(self):
        # GH 19976
        result = np.all(DataFrame(columns=["a", "b"])).item()
        assert result is True

        result = np.any(DataFrame(columns=["a", "b"])).item()
        assert result is False
//  [python]
def test_multinomial_vector(self):
        rng_R = random_state_type()
        n = tensor.lvector()
        pvals = tensor.matrix()
        post_r, out = multinomial(rng_R, n=n, pvals=pvals)
        assert out.ndim == 2
        f = compile.function([rng_R, n, pvals], [post_r, out],
                             accept_inplace=True)

        n_val = [1, 2, 3]
        pvals_val = [[.1, .9], [.2, .8], [.3, .7]]
        pvals_val = numpy.asarray(pvals_val, dtype=config.floatX)
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, n_val, pvals_val)
        numpy_val0 = numpy.asarray([numpy_rng.multinomial(n=nv, pvals=pv)
            for nv, pv in zip(n_val, pvals_val)])
        assert n.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, n_val[:-1], pvals_val[:-1])
        numpy_val1 = numpy.asarray([numpy_rng.multinomial(n=nv, pvals=pv)
            for nv, pv in zip(n_val[:-1], pvals_val[:-1])])
        assert n.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, n, pvals],
                multinomial(rng_R, n=n, pvals=pvals, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, n_val, pvals_val)
        numpy_val2 = numpy.asarray([numpy_rng.multinomial(n=nv, pvals=pv)
            for nv, pv in zip(n_val, pvals_val)])
        assert n.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, n_val[:-1], pvals_val[:-1])
//  [python]
def test_multinomial_vector(self):
        rng_R = random_state_type()
        n = tensor.lvector()
        pvals = tensor.matrix()
        post_r, out = multinomial(rng_R, n=n, pvals=pvals)
        assert out.ndim == 2
        f = compile.function([rng_R, n, pvals], [post_r, out],
                             accept_inplace=True)

        n_val = [1, 2, 3]
        pvals_val = [[.1, .9], [.2, .8], [.3, .7]]
        pvals_val = numpy.asarray(pvals_val, dtype=config.floatX)
        rng = numpy.random.RandomState(utt.fetch_seed())
        numpy_rng = numpy.random.RandomState(utt.fetch_seed())

        # Arguments of size (3,)
        rng0, val0 = f(rng, n_val, pvals_val)
        numpy_val0 = numpy.asarray([numpy_rng.multinomial(n=nv, pvals=pv)
            for nv, pv in zip(n_val, pvals_val)])
        assert numpy.all(val0 == numpy_val0)

        # arguments of size (2,)
        rng1, val1 = f(rng0, n_val[:-1], pvals_val[:-1])
        numpy_val1 = numpy.asarray([numpy_rng.multinomial(n=nv, pvals=pv)
            for nv, pv in zip(n_val[:-1], pvals_val[:-1])])
        assert numpy.all(val1 == numpy_val1)

        # Specifying the size explicitly
        g = compile.function([rng_R, n, pvals],
                multinomial(rng_R, n=n, pvals=pvals, size=(3,)),
                accept_inplace=True)
        rng2, val2 = g(rng1, n_val, pvals_val)
        numpy_val2 = numpy.asarray([numpy_rng.multinomial(n=nv, pvals=pv)
            for nv, pv in zip(n_val, pvals_val)])
        assert numpy.all(val2 == numpy_val2)
        self.assertRaises(ValueError, g, rng2, n_val[:-1], pvals_val[:-1])
//  [python]
def testEigHandlesNanInputs(self, shape, dtype, compute_left_eigenvectors,
                              compute_right_eigenvectors):
    """Verifies that `eig` fails gracefully if given non-finite inputs."""
    a = jnp.full(shape, jnp.nan, dtype)
    results = lax.linalg.eig(
        a, compute_left_eigenvectors=compute_left_eigenvectors,
        compute_right_eigenvectors=compute_right_eigenvectors)
    for result in results:
      self.assertTrue(np.all(np.isnan(result)))
//  [python]
def testQrBatching(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    args = rng(shape, jnp.float32)
    qs, rs = vmap(jsp.linalg.qr)(args)
    self.assertTrue(np.all(np.linalg.norm(args - np.matmul(qs, rs)) < 1e-3))
//  [python]
def testEigBatching(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    shape = (10,) + shape
    args = rng(shape, dtype)
    ws, vs = vmap(jnp.linalg.eig)(args)
    self.assertTrue(np.all(np.linalg.norm(
        np.matmul(args, vs) - ws[..., None, :] * vs) < 1e-3))
//  [python]
def test_numpy_frames(small_robot_mp4_vid):
    clip_frames = small_robot_mp4_vid.get_frames((3, 7, 9))
    np_vid = Video.from_numpy(clip_frames)

    assert np.alltrue(np.equal(np_vid.get_frame(1), small_robot_mp4_vid.get_frame(7)))
//  [python]
def test_numpy_frames(small_robot_mp4_vid):
    clip_frames = small_robot_mp4_vid.get_frames((3, 7, 9))
    np_vid = Video.from_numpy(clip_frames)

    assert np.all(np.equal(np_vid.get_frame(1), small_robot_mp4_vid.get_frame(7)))
//  [python]
def tst_basic(self, x, T, mask, val):
        np.putmask(x, mask, val)
        assert_(n.all(x[mask] == T(val)))
        assert_(x.dtype == T)
//  [python]
def tst_basic(self, x, T, mask, val):
        np.putmask(x, mask, val)
        assert_(np.all(x[mask] == T(val)))
        assert_(x.dtype == T)
//  [python]
def testEig(self, shape, dtype, compute_left_eigenvectors,
              compute_right_eigenvectors):
    rng = jtu.rand_default(self.rng())
    n = shape[-1]
    args_maker = lambda: [rng(shape, dtype)]

    # Norm, adjusted for dimension and type.
    def norm(x):
      norm = np.linalg.norm(x, axis=(-2, -1))
      return norm / ((n + 1) * jnp.finfo(dtype).eps)

    def check_right_eigenvectors(a, w, vr):
      self.assertTrue(
        n.all(norm(np.matmul(a, vr) - w[..., None, :] * vr) < 100))

    def check_left_eigenvectors(a, w, vl):
      rank = len(a.shape)
      aH = jnp.conj(a.transpose(list(range(rank - 2)) + [rank - 1, rank - 2]))
      wC = jnp.conj(w)
      check_right_eigenvectors(aH, wC, vl)

    a, = args_maker()
    results = lax.linalg.eig(
        a, compute_left_eigenvectors=compute_left_eigenvectors,
        compute_right_eigenvectors=compute_right_eigenvectors)
    w = results[0]

    if compute_left_eigenvectors:
      check_left_eigenvectors(a, w, results[1])
    if compute_right_eigenvectors:
      check_right_eigenvectors(a, w, results[1 + compute_left_eigenvectors])

    self._CompileAndCheck(partial(jnp.linalg.eig), args_maker,
                          rtol=1e-3)
//  [python]
def testEig(self, shape, dtype, compute_left_eigenvectors,
              compute_right_eigenvectors):
    rng = jtu.rand_default(self.rng())
    n = shape[-1]
    args_maker = lambda: [rng(shape, dtype)]

    # Norm, adjusted for dimension and type.
    def norm(x):
      norm = np.linalg.norm(x, axis=(-2, -1))
      return norm / ((n + 1) * jnp.finfo(dtype).eps)

    def check_right_eigenvectors(a, w, vr):
      self.assertTrue(
        np.all(norm(np.matmul(a, vr) - w[..., None, :] * vr) < 100))

    def check_left_eigenvectors(a, w, vl):
      rank = len(a.shape)
      aH = jnp.conj(a.transpose(list(range(rank - 2)) + [rank - 1, rank - 2]))
      wC = jnp.conj(w)
      check_right_eigenvectors(aH, wC, vl)

    a, = args_maker()
    results = lax.linalg.eig(
        a, compute_left_eigenvectors=compute_left_eigenvectors,
        compute_right_eigenvectors=compute_right_eigenvectors)
    w = results[0]

    if compute_left_eigenvectors:
      check_left_eigenvectors(a, w, results[1])
    if compute_right_eigenvectors:
      check_right_eigenvectors(a, w, results[1 + compute_left_eigenvectors])

    self._CompileAndCheck(partial(jnp.linalg.eig), args_maker,
                          rtol=1e-3)
//  [python]
def numpy_scalar(data):
    """ Return a scalar stored in a numpy ndarray.

    Raises
    ------
     NotScalarConstantError
        If the numpy ndarray is not a scalar.

    """

    # handle case where data is numpy.array([])
    if (data.ndim > 0 and
        (len(data.shape) == 0 or
         __builtins__['max'](data.shape) == 0)):
        assert numpy.all(numpy.array([]) == data)
        raise EmptyConstantError()
    try:
        numpy.complex(data)  # works for all numeric scalars
        return data
    except Exception:
        raise NotScalarConstantError(
            'v.data is non-numeric, non-scalar, or has more than one'
            ' unique value', data)
//  [python]
def test_any_all_object(self):
        # https://github.com/pandas-dev/pandas/issues/19976
        result = np.all(DataFrame(columns=['a', 'b'])).item()
        assert result is True

        result = np.any(DataFrame(columns=['a', 'b'])).item()
        assert result is False
//  [python]
def nansum_missing(array, axis):
    """nansum where all-NaN values remain NaNs.

    Note: In NumPy <=1.9 NaN is returned for slices that are
    all NaN, while later versions return 0. This function emulates
    the older behavior, which allows using NaN as a missing value
    indicator.

    Parameters
    ----------
    array: Array to sum over
    axis:  Axis to sum over
    """
    T = list(range(array.ndim))
    T.remove(axis)
    T.insert(0, axis)
    array = array.transpose(T)
    missing_vals = np.isnan(array)
    all_empty = np.alltrue(missing_vals, axis=0)
    set_to_zero = missing_vals & ~all_empty
    return np.where(set_to_zero, 0, array).sum(axis=0)
//  [python]
def nansum_missing(array, axis):
    """nansum where all-NaN values remain NaNs.

    Note: In NumPy <=1.9 NaN is returned for slices that are
    all NaN, while later versions return 0. This function emulates
    the older behavior, which allows using NaN as a missing value
    indicator.

    Parameters
    ----------
    array: Array to sum over
    axis:  Axis to sum over
    """
    T = list(range(array.ndim))
    T.remove(axis)
    T.insert(0, axis)
    array = array.transpose(T)
    missing_vals = np.isnan(array)
    all_empty = np.all(missing_vals, axis=0)
    set_to_zero = missing_vals & ~all_empty
    return np.where(set_to_zero, 0, array).sum(axis=0)
//  [python]
def test_flat(self):
        # Test that flat can return all types of items [#4585, #4615]
        # test simple access
        test = masked_array(np.matrix([[1, 2, 3]]), mask=[0, 0, 1])
        assert_equal(test.flat[1], 2)
        assert_equal(test.flat[2], masked)
        self.assertTrue(np.all(test.flat[0:2] == test[0, 0:2]))
        # Test flat on masked_matrices
        test = masked_array(np.matrix([[1, 2, 3]]), mask=[0, 0, 1])
        test.flat = masked_array([3, 2, 1], mask=[1, 0, 0])
        control = masked_array(np.matrix([[3, 2, 1]]), mask=[1, 0, 0])
        assert_equal(test, control)
        # Test setting
        test = masked_array(np.matrix([[1, 2, 3]]), mask=[0, 0, 1])
        testflat = test.flat
        testflat[:] = testflat[[2, 1, 0]]
        assert_equal(test, control)
        testflat[0] = 9
        assert_equal(test[0, 0], 9)
        # test 2-D record array
        # ... on structured array w/ masked records
        x = array([[(1, 1.1, 'one'), (2, 2.2, 'two'), (3, 3.3, 'thr')],
                   [(4, 4.4, 'fou'), (5, 5.5, 'fiv'), (6, 6.6, 'six')]],
                  dtype=[('a', int), ('b', float), ('c', '|S8')])
        x['a'][0, 1] = masked
        x['b'][1, 0] = masked
        x['c'][0, 2] = masked
        x[-1, -1] = masked
        xflat = x.flat
        assert_equal(xflat[0], x[0, 0])
        assert_equal(xflat[1], x[0, 1])
        assert_equal(xflat[2], x[0, 2])
        assert_equal(xflat[:3], x[0])
        assert_equal(xflat[3], x[1, 0])
        assert_equal(xflat[4], x[1, 1])
        assert_equal(xflat[5], x[1, 2])
        assert_equal(xflat[3:], x[1])
        assert_equal(xflat[-1], x[-1, -1])
        i = 0
        j = 0
        for xf in xflat:
            assert_equal(xf, x[j, i])
            i += 1
            if i >= x.shape[-1]:
                i = 0
                j += 1
        # test that matrices keep the correct shape (#4615)
        a = masked_array(np.matrix(np.eye(2)), mask=0)
        b = a.flat
        b01 = b[:2]
        assert_equal(b01.data, array([[1., 0.]]))
        assert_equal(b01.mask, array([[False, False]]))
//  [python]
def test_fromiter_bytes(self):
        # Ticket #1058
        a = np.fromiter(list(range(10)), dtype='b')
        b = np.fromiter(list(range(10)), dtype='B')
        assert_(a == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])).all()
        assert_(b == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])).all()
//  [python]
def test_fromiter_comparison(self):
        a = np.fromiter(list(range(10)), dtype='b')
        b = np.fromiter(list(range(10)), dtype='B')
        assert_(a == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]).all())
        assert_(b == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]).all())
//  [python]
def _build_y_x(da, tolerance):
    """Build y/x coordinates from DataArray's latitude/longitude coordinates."""
    # Initial sanity checks
    latitude, longitude = da.metpy.coordinates('latitude', 'longitude')
    if latitude.dims != longitude.dims:
        raise ValueError('Latitude and longitude must have same dimensionality')
    elif latitude.ndim != 2:
        raise ValueError('To build 1D y/x coordinates via assign_y_x, latitude/longitude '
                         'must be 2D')

    # Convert to projected y/x
    xxyy = np.stack(da.metpy.pyproj_proj(
        longitude.values,
        latitude.values,
        inverse=False,
        radians=False
    ), axis=-1)

    # Handle tolerance
    tolerance = 1 if tolerance is None else tolerance.m_as('m')

    # If within tolerance, take median to collapse to 1D
    try:
        y_dim = latitude.metpy.find_axis_number('y')
        x_dim = latitude.metpy.find_axis_number('x')
    except AttributeError:
        _warnings.warn('y and x dimensions unable to be identified. Assuming [..., y, x] '
                       'dimension order.')
        y_dim, x_dim = 0, 1
    if (np.alltrue(np.ptp(xxyy[..., 0], axis=y_dim) < tolerance)
            and np.alltrue(np.ptp(xxyy[..., 1], axis=x_dim) < tolerance)):
        x = np.median(xxyy[..., 0], axis=y_dim)
        y = np.median(xxyy[..., 1], axis=x_dim)
        x = xr.DataArray(x, name=latitude.dims[x_dim], dims=(latitude.dims[x_dim],),
                         coords={latitude.dims[x_dim]: x},
                         attrs={'units': 'meter', 'standard_name': 'projection_x_coordinate'})
        y = xr.DataArray(y, name=latitude.dims[y_dim], dims=(latitude.dims[y_dim],),
                         coords={latitude.dims[y_dim]: y},
                         attrs={'units': 'meter', 'standard_name': 'projection_y_coordinate'})
        return y, x
    else:
        raise ValueError('Projected y and x coordinates cannot be collapsed to 1D within '
                         'tolerance. Verify that your latitude and longitude coordinates '
                         'correspond to your CRS coordinate.')
//  [python]
def _build_y_x(da, tolerance):
    """Build y/x coordinates from DataArray's latitude/longitude coordinates."""
    # Initial sanity checks
    latitude, longitude = da.metpy.coordinates('latitude', 'longitude')
    if latitude.dims != longitude.dims:
        raise ValueError('Latitude and longitude must have same dimensionality')
    elif latitude.ndim != 2:
        raise ValueError('To build 1D y/x coordinates via assign_y_x, latitude/longitude '
                         'must be 2D')

    # Convert to projected y/x
    xxyy = np.stack(da.metpy.pyproj_proj(
        longitude.values,
        latitude.values,
        inverse=False,
        radians=False
    ), axis=-1)

    # Handle tolerance
    tolerance = 1 if tolerance is None else tolerance.m_as('m')

    # If within tolerance, take median to collapse to 1D
    try:
        y_dim = latitude.metpy.find_axis_number('y')
        x_dim = latitude.metpy.find_axis_number('x')
    except AttributeError:
        _warnings.warn('y and x dimensions unable to be identified. Assuming [..., y, x] '
                       'dimension order.')
        y_dim, x_dim = 0, 1
    if (np.all(np.ptp(xxyy[..., 0], axis=y_dim) < tolerance)
            and np.all(np.ptp(xxyy[..., 1], axis=x_dim) < tolerance)):
        x = np.median(xxyy[..., 0], axis=y_dim)
        y = np.median(xxyy[..., 1], axis=x_dim)
        x = xr.DataArray(x, name=latitude.dims[x_dim], dims=(latitude.dims[x_dim],),
                         coords={latitude.dims[x_dim]: x},
                         attrs={'units': 'meter', 'standard_name': 'projection_x_coordinate'})
        y = xr.DataArray(y, name=latitude.dims[y_dim], dims=(latitude.dims[y_dim],),
                         coords={latitude.dims[y_dim]: y},
                         attrs={'units': 'meter', 'standard_name': 'projection_y_coordinate'})
        return y, x
    else:
        raise ValueError('Projected y and x coordinates cannot be collapsed to 1D within '
                         'tolerance. Verify that your latitude and longitude coordinates '
                         'correspond to your CRS coordinate.')
//  [python]
def test_numpy_rw_masked_array_2():

    ds = gdal.Open("data/test3_with_mask_8bit.tif")

    band = ds.GetRasterBand(1)

    arr = band.ReadAsArray()
    mask = band.GetMaskBand().ReadAsArray()

    masked_arr = band.ReadAsMaskedArray()

    assert not numpy.any(masked_arr.mask[mask == 255])
    assert n.all(masked_arr.mask[mask != 255])

    assert masked_arr.sum() == arr[mask == 255].sum()
//  [python]
def test_numpy_rw_masked_array_2():

    ds = gdal.Open("data/test3_with_mask_8bit.tif")

    band = ds.GetRasterBand(1)

    arr = band.ReadAsArray()
    mask = band.GetMaskBand().ReadAsArray()

    masked_arr = band.ReadAsMaskedArray()

    assert not numpy.any(masked_arr.mask[mask == 255])
    assert numpy.all(masked_arr.mask[mask != 255])

    assert masked_arr.sum() == arr[mask == 255].sum()
//  [python]
def test_numpy_rw_18():

    img = numpy.random.randint(0, 255, size=(256, 200, 3)).astype("uint8")
    ds = gdal_array.OpenArray(img, interleave="pixel")
    assert ds is not None, "Failed to open memory array as dataset."

    bnd1 = ds.GetRasterBand(1).ReadAsArray()
    bnd2 = ds.GetRasterBand(2).ReadAsArray()
    bnd3 = ds.GetRasterBand(3).ReadAsArray()

    res = numpy.dstack((bnd1, bnd2, bnd3))
    assert n.all(img == res)

    res = ds.ReadAsArray(interleave="pixel")
    assert n.all(img == res)

    res = numpy.zeros([256, 200, 3])
    ds.ReadAsArray(buf_obj=res, interleave="pixel")
    assert n.all(img == res)
//  [python]
def test_numpy_rw_18():

    img = numpy.random.randint(0, 255, size=(256, 200, 3)).astype("uint8")
    ds = gdal_array.OpenArray(img, interleave="pixel")
    assert ds is not None, "Failed to open memory array as dataset."

    bnd1 = ds.GetRasterBand(1).ReadAsArray()
    bnd2 = ds.GetRasterBand(2).ReadAsArray()
    bnd3 = ds.GetRasterBand(3).ReadAsArray()

    res = numpy.dstack((bnd1, bnd2, bnd3))
    assert numpy.all(img == res)

    res = ds.ReadAsArray(interleave="pixel")
    assert numpy.all(img == res)

    res = numpy.zeros([256, 200, 3])
    ds.ReadAsArray(buf_obj=res, interleave="pixel")
    assert numpy.all(img == res)
//  [python]
def test_categorical_likelihood(
    num_cats: int, cat_probs: np.ndarray, hybridize: bool
):
    """
    Test to check that maximizing the likelihood recovers the parameters
    """
    cat_prob = mx.nd.array(cat_probs)
    cat_probs = mx.nd.zeros((NUM_SAMPLES, num_cats)) + cat_prob

    distr = Categorical(cat_probs.log())
    samples = distr.sample()

    cat_prob_init = mx.nd.random_uniform(1 - TOL, 1 + TOL, num_cats) * cat_prob
    cat_prob_init = cat_prob_init / cat_prob_init.sum()

    init_biases = [cat_prob_init]

    cat_log_prob_hat = maximum_likelihood_estimate_sgd(
        CategoricalOutput(num_cats),
        samples,
        init_biases=init_biases,
        hybridize=hybridize,
        learning_rate=PositiveFloat(0.05),
        num_epochs=PositiveInt(25),
    )
    cat_prob_hat = np.exp(cat_log_prob_hat)

    prob_deviation = np.abs(cat_prob_hat - cat_prob.asnumpy()).flatten()
    tolerance = (TOL * cat_prob.asnumpy()).flatten()

    assert np.all(
        np.less(prob_deviation, tolerance)
    ), f"cat_prob did not match: cat_prob = {cat_prob}, cat_prob_hat = {cat_prob_hat}"
//  [python]
def testRngUniformS32(self):
    lo, hi = 2, 4
    shape = (2, 3)
    c = self._NewComputation()
    c.RngUniform(
        c.Constant(NumpyArrayS32(lo)),
        c.Constant(NumpyArrayS32(hi)),
        dims=shape)
    result = xla_client.execute_with_python_values(c.Build().Compile())
    # since the result is random, we just check shape, integrality, and range
    self.assertEqual(result.shape, shape)
    self.assertEqual(result.dtype, np.int32)
    self.assertTrue(numpy.all(lo <= result))
    self.assertTrue(numpy.all(result < hi))
//  [python]
def testRngUniformS32(self):
    lo, hi = 2, 4
    shape = (2, 3)
    c = self._NewComputation()
    c.RngUniform(
        c.Constant(NumpyArrayS32(lo)),
        c.Constant(NumpyArrayS32(hi)),
        dims=shape)
    result = xla_client.execute_with_python_values(c.Build().Compile())
    # since the result is random, we just check shape, integrality, and range
    self.assertEqual(result.shape, shape)
    self.assertEqual(result.dtype, np.int32)
    self.assertTrue(np.all(lo <= result))
    self.assertTrue(np.all(result < hi))