jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def setup_data(p, test_set=False):
    if p.dataset in ['cifar10','mnist']:
        dataset_class, training_set_size = {
            'cifar10': (CIFAR10, 40000),
            'mnist': (MNIST, 50000),
        }[p.dataset]
    else:
        from fuel.datasets import H5PYDataset
        from fuel.utils import find_in_data_path
        from functools import partial
        fn=p.dataset
        fn=os.path.join(fn, fn + '.hdf5')
        def dataset_class(which_sets):
            return H5PYDataset(file_or_path=find_in_data_path(fn),
                               which_sets=which_sets,
                               load_in_memory=True)
        training_set_size = None

    train_set = dataset_class(["train"])

    # Allow overriding the default from command line
    if p.get('unlabeled_samples') is not None and p.unlabeled_samples >= 0:
        training_set_size = p.unlabeled_samples
    elif training_set_size is None:
        training_set_size = train_set.num_examples

    # Make sure the MNIST data is in right format
    if p.dataset == 'mnist':
        d = train_set.data_sources[train_set.sources.index('features')]
        assert numpy.all(d <= 1.0) and numpy.all(d >= 0.0), \
            'Make sure data is in float format and in range 0 to 1'

    # Take all indices and permutate them
    all_ind = numpy.arange(train_set.num_examples)
    if p.get('dseed'):
        rng = numpy.random.RandomState(seed=p.dseed)
        rng.shuffle(all_ind)

    d = AttributeDict()

    # Choose the training set
    d.train = train_set
    d.train_ind = all_ind[:training_set_size]

    # Then choose validation set from the remaining indices
    d.valid = train_set
    d.valid_ind = numpy.setdiff1d(all_ind, d.train_ind)[:p.valid_set_size]
    logger.info('Using %d examples for validation' % len(d.valid_ind))

    # Only touch test data if requested
    if test_set:
        d.test = dataset_class(["test"])
        d.test_ind = numpy.arange(d.test.num_examples)

    # Setup optional whitening, only used for Cifar-10
    in_dim = train_set.data_sources[train_set.sources.index('features')].shape[1:]
    if len(in_dim) > 1 and p.whiten_zca > 0:
        assert numpy.product(in_dim) == p.whiten_zca, \
            'Need %d whitening dimensions, not %d' % (numpy.product(in_dim),
                                                      p.whiten_zca)
    cnorm = ContrastNorm(p.contrast_norm) if p.contrast_norm != 0 else None

    def get_data(d, i):
        data = d.get_data(request=i)[d.sources.index('features')]
        # Fuel provides Cifar in uint8, convert to float32
        data = numpy.require(data, dtype=numpy.float32)
        return data if cnorm is None else cnorm.apply(data)

    if p.whiten_zca > 0:
        logger.info('Whitening using %d ZCA components' % p.whiten_zca)
        whiten = ZCA()
        whiten.fit(p.whiten_zca, get_data(d.train, d.train_ind))
    else:
        whiten = None

    return in_dim, d, whiten, cnorm
//  [python]
def get_duplicates(M):
    res = []
    I = np.lexsort([M[:, i] for i in reversed(range(0, M.shape[1]))])
    S = M[I, :]

    i = 0

    while i < S.shape[0] - 1:
        l = []
        while n.all(S[i, :] == S[i + 1, :]):
            l.append(I[i])
            i += 1
        if len(l) > 0:
            l.append(I[i])
            res.append(l)
        i += 1

    return res
//  [python]
def get_duplicates(M):
    res = []
    I = np.lexsort([M[:, i] for i in reversed(range(0, M.shape[1]))])
    S = M[I, :]

    i = 0

    while i < S.shape[0] - 1:
        l = []
        while np.all(S[i, :] == S[i + 1, :]):
            l.append(I[i])
            i += 1
        if len(l) > 0:
            l.append(I[i])
            res.append(l)
        i += 1

    return res
//  [python]
def add_calculate(num, denum, aslist=False, out_type=None):
    # TODO: make sure that this function and mul_calculate are similar
    if out_type is None:
        zero = 0.0
    else:
        zero = theano._asarray(0, dtype=out_type.dtype)
    # zero = 0.0 if out_type is None else theano._asarray(0,
    # dtype=out_type.dtype)
    v = reduce(N.add, num, zero) - reduce(N.add, denum, zero)
    if aslist:
        if N.all(v == 0):
            return []
        else:
            return [v]
    return v
//  [python]
def mul_calculate(num, denum, aslist=False, out_type=None):
    if not num and not denum:
        # Smallest 1 possible.
        if aslist:
            return []
        else:
            return N.int8(1)

    # Make sure we do not accidently upcast data types.
    if out_type is None:
        out_dtype = scalar.upcast(*[v.dtype for v in (num + denum)])
    else:
        out_dtype = out_type.dtype
    one = theano._asarray(1, dtype=out_dtype)

    v = reduce(N.multiply, num, one) / reduce(N.multiply, denum, one)
    if aslist:
        if N.all(v == 1):
            return []
        else:
            return [v]
    return v
//  [python]
def get_degeneracies ( E_k, bnd ):

  all_degen = []

  E_k_round = np.around(E_k, decimals=5)
  
  for ispin in range(E_k_round.shape[2]):

    by_spin = []
    for ik in range(E_k_round.shape[0]):
      by_kp = []
      eV = np.unique(E_k_round[ik,:,ispin][:-1][np.isclose(E_k_round[ik,:,ispin][1:],
                                                           E_k_round[ik,:,ispin][:-1],
                                                           atol=1.e-5)])

      for i in range(len(eV)):
        inds = np.where(np.isclose(E_k_round[ik,:,ispin],eV[i],atol=1.e-5))[0]

        if len(inds) > 1 and np.all(inds < bnd):                    
          by_kp.append(inds)

      by_spin.append(by_kp)

    all_degen.append(by_spin)

  return all_degen
//  [python]
def test_indexing_single_index(self):
        index = 12
        ds = self._dataset()

        result = ds[index]
        self.assertEqual(1, len(result))
        self.assertTrue(isinstance(result, self.dataset_view_class))

        expected_x = [ds.x[index]]
        self.assertTrue(np.all(expected_x == result.x))
//  [python]
def test_indexing_list_index(self):
        index = [1, 2, 7, 13]
        ds = self._dataset()

        result = ds[index]
        self.assertEqual(4, len(result))
        self.assertTrue(isinstance(result, self.dataset_view_class))

        expected_x = [ds.x[i] for i in index]
        self.assertTrue(np.all(expected_x == result.x))
//  [python]
def test_indexing_slicing(self):
        index = np.s_[4:11]
        ds = self._dataset()

        result = ds[index]
        self.assertEqual(7, len(result))
        self.assertTrue(isinstance(result, self.dataset_view_class))

        expected_x = [ds.x[i] for i in np.arange(len(ds))[index]]
        self.assertTrue(np.all(expected_x == result.x))
//  [python]
def test_clone(self):
        ds = self._random_data(num_samples=self.NUM_SAMPLES)
        ds_cloned = ds.clone()

        self.assertTrue(np.all([
            torch.equal(t, t_cloned)
            for t, t_cloned in zip(ds.x, ds_cloned.x)
        ]))
        if self.multi_label:
            assert_csr_matrix_equal(ds.y, ds_cloned.y)
        else:
            assert_array_equal(ds.y, ds_cloned.y)
        if self.target_labels == 'explicit':
            assert_array_equal(ds.target_labels, ds_cloned.target_labels)

        # test propagation of target_labels setting
        if self.target_labels == 'explicit':
            self.assertFalse(ds.track_target_labels)
            self.assertFalse(ds_cloned.track_target_labels)
            assert_array_equal(ds.target_labels, ds_cloned.target_labels)
        else:
            self.assertTrue(ds.track_target_labels)
            self.assertTrue(ds_cloned.track_target_labels)

        # mutability test
        ds_cloned.x[0][0] += 1
        assert_list_of_tensors_not_equal(self, ds.x, ds_cloned.x)

        if self.multi_label:
            y_tmp = ds_cloned.y.todense()
            y_tmp[2:10] = 0
            ds_cloned.y = csr_matrix(y_tmp)
            assert_array_not_equal(ds.y.indices, ds_cloned.y.indices)
        else:
            ds_cloned = increase_dense_labels_safe(ds_cloned)
            assert_array_not_equal(ds.y, ds_cloned.y)

        ds_cloned.target_labels = np.arange(10)
        assert_array_not_equal(ds.target_labels, ds_cloned.target_labels)
//  [python]
def contains(grid, points):
    return numpy.logical_and(
        numpy.alltrue(points >= grid.get_minimum_position(), axis=len(points.shape)-1),
        numpy.alltrue(points < grid.get_maximum_position(), axis=len(points.shape)-1)
    )
//  [python]
def contains(grid, points):
    return numpy.logical_and(
        numpy.all(points >= grid.get_minimum_position(), axis=len(points.shape)-1),
        numpy.all(points < grid.get_maximum_position(), axis=len(points.shape)-1)
    )
//  [python]
def _clone_test(self, ds_view):
        ds_cloned = ds_view.clone()
        self.assertTrue(isinstance(ds_cloned, self.dataset_class))

        self.assertTrue(np.all([
            torch.equal(t, t_cloned)
            for t, t_cloned in zip(ds_view.x, ds_cloned.x)
        ]))
        if self.multi_label:
            if self.target_labels == 'explicit':
                assert_csr_matrix_equal(ds_view.y, ds_cloned.y)
        else:
            assert_array_equal(ds_view.y, ds_cloned.y)
        assert_array_equal(ds_view.target_labels, ds_cloned.target_labels)

        # test propagation of target_labels setting
        if self.target_labels == 'explicit':
            self.assertFalse(ds_view.dataset.track_target_labels)
            self.assertFalse(ds_cloned.track_target_labels)
            assert_array_equal(ds_view.target_labels, ds_cloned.target_labels)
        else:
            self.assertTrue(ds_view.dataset.track_target_labels)
            self.assertTrue(ds_cloned.track_target_labels)

        # mutability test
        ds_cloned.x[0][0] += 1
        assert_list_of_tensors_not_equal(self, ds_view.x, ds_cloned.x)

        if self.multi_label:
            y_tmp = ds_cloned.y.todense()
            y_tmp = (y_tmp + 1) % 2
            ds_cloned.y = csr_matrix(y_tmp, shape=ds_view.y.shape)
            try:
                assert_csr_matrix_not_equal(ds_view.y, ds_cloned.y)
            except (AssertionError, ValueError):
                assert_csr_matrix_not_equal(ds_view.y, ds_cloned.y)
        else:
            ds_cloned = increase_dense_labels_safe(ds_cloned)
            assert_array_not_equal(ds_view.y, ds_cloned.y)

        ds_cloned.target_labels = np.arange(10)
        assert_array_not_equal(ds_view.target_labels, ds_cloned.target_labels)
//  [python]
def test_eigh(self, compute_v, shape):
    np.random.seed(42)
    inputs = np.random.normal(size=shape).astype(np.float32)

    @tf.function(jit_compile=True)
    def eigh(x):
      return tf.raw_ops.SelfAdjointEigV2(input=x, compute_v=compute_v)

    tf_w, tf_v = eigh(tf.constant(inputs))

    jax_eigh = tf2jax.convert_functional(eigh, np.zeros_like(inputs))
    jax_w, jax_v = self.variant(jax_eigh)(inputs)

    self.assertTrue(np.all(tf_w[..., :-1] <= tf_w[..., 1:]))
    self.assertTrue(np.all(jax_w[..., :-1] <= jax_w[..., 1:]))

    # Check eigenvalues so because ordering is not consistent even after sorting
    for idx in itertools.product(*map(range, shape[:-2])):
      w_jax, w_tf = jax_w[idx], tf_w[idx]
      for i in range(shape[-1]):
        self.assertAllClose(min(abs(w_jax - w_tf[i])), 0., atol=1e-5)
        self.assertAllClose(min(abs(w_jax[i] - w_tf)), 0., atol=1e-5)

    if compute_v:
      tols = dict(atol=2e-5) if jax.default_backend() in "tpu" else {}
      sym_inputs = np.tril(inputs) + np.swapaxes(np.tril(inputs, -1), -2, -1)
      self.assertAllClose(
          np.matmul(sym_inputs, tf_v), tf_w[..., None, :] * tf_v, **tols)
      self.assertAllClose(
          np.matmul(sym_inputs, jax_v), jax_w[..., None, :] * jax_v, **tols)
//  [python]
def is_real_array(array):
    '''
    Check if a value is a list or array of real scalars by aggregating several
    numpy checks.

    Parameters
    ----------
    array: any type
        The parameter to check

    Returns
    ------
    check : bool
        True if array is a list or array of a real scalars, False
        otherwise.

    '''

    if not (type(array) is list or type(array) is np.ndarray):
        return False
    else:
        if (not all([np.isreal(x) for x in array]) or
                not np.isrealobj(array) or
                not np.asarray(list(map(np.isscalar, array))).all()):
            return False
        else:
            return True
//  [python]
def is_real_array(array):
    '''
    Check if a value is a list or array of real scalars by aggregating several
    numpy checks.

    Parameters
    ----------
    array: any type
        The parameter to check

    Returns
    ------
    check : bool
        True if ```array``` is a list or array of a real scalars, False
        otherwise.

    '''

    if not (type(array) is list or type(array) is np.ndarray):
        return False
    else:
        if (not np.all([np.isreal(x) for x in array]) or
                not np.isrealobj(array) or
                not np.asarray(list(map(np.isscalar, array))).all()):
            return False
        else:
            return True
//  [python]
def testPointsInside(self,points):
        """tests if a point is inside the convex mesh of the robot.

        Returns an array the same length as points that specifies whether the point is in or not. This method is not meant to be optimized (use C++ for that).
        """
        with self.env:
            inside = zeros(len(points),bool)
            leftinds = arange(len(points))
            leftpoints = points[:]
            for ilink,link in enumerate(self.robot.GetLinks()):
                if len(leftinds) == 0:
                    break
                for ig,geom in enumerate(link.GetGeometries()):
                    T = dot(link.GetTransform(),geom.GetTransform())
                    localpoints = transformInversePoints(T,leftpoints)
                    cdhulls = [cdhull for ig2,cdhull in self.linkgeometry[ilink] if ig2==ig]
                    if len(cdhulls) > 0:
                        for ihull,hull in enumerate(cdhulls[0]):
                            insideinds = n.all(dot(localpoints,transpose(hull[2][:,0:3]))+tile(hull[2][:,3],(len(localpoints),1))<=0,1)
                            inside[leftinds[flatnonzero(insideinds)]] = True
                            outsideinds = flatnonzero(insideinds==0)
                            if len(outsideinds) > 0:
                                leftpoints = leftpoints[outsideinds]
                                leftinds = leftinds[outsideinds]
                                localpoints = localpoints[outsideinds]
                            if len(leftinds) == 0:
                                break
                        continue
                    elif geom.GetType() == KinBody.Link.GeomType.Box:
                        insideinds = n.all(n.less_equal(abs(localpoints), geom.GetBoxExtents()) ,1)
                    elif geom.GetType() == KinBody.Link.GeomType.Sphere:
                        insideinds = n.less_equal(sum(localpoints**2,1), geom.GetSphereRadius()**2)
                    elif geom.GetType() == KinBody.Link.GeomType.Cylinder:
                        insideinds = n.less_equal(abs(localpoints[:,1]), 0.5*geom.GetCylinderHeight()) and n.less_equal(localpoint[:,0]**2+localpoint[:2]**2, geom.GetCylinderRadius()**2)
                    else:
                        continue
                    inside[leftinds[flatnonzero(insideinds)]] = True
                    outsideinds = flatnonzero(insideinds==0)
                    if len(outsideinds) > 0:
                        leftpoints = leftpoints[outsideinds]
                        leftinds = leftinds[outsideinds]
                    if len(leftinds) == 0:
                        break
        return inside
//  [python]
def testPointsInside(self,points):
        """tests if a point is inside the convex mesh of the robot.

        Returns an array the same length as points that specifies whether the point is in or not. This method is not meant to be optimized (use C++ for that).
        """
        with self.env:
            inside = zeros(len(points),bool)
            leftinds = arange(len(points))
            leftpoints = points[:]
            for ilink,link in enumerate(self.robot.GetLinks()):
                if len(leftinds) == 0:
                    break
                for ig,geom in enumerate(link.GetGeometries()):
                    T = dot(link.GetTransform(),geom.GetTransform())
                    localpoints = transformInversePoints(T,leftpoints)
                    cdhulls = [cdhull for ig2,cdhull in self.linkgeometry[ilink] if ig2==ig]
                    if len(cdhulls) > 0:
                        for ihull,hull in enumerate(cdhulls[0]):
                            insideinds = numpy.all(dot(localpoints,transpose(hull[2][:,0:3]))+tile(hull[2][:,3],(len(localpoints),1))<=0,1)
                            inside[leftinds[flatnonzero(insideinds)]] = True
                            outsideinds = flatnonzero(insideinds==0)
                            if len(outsideinds) > 0:
                                leftpoints = leftpoints[outsideinds]
                                leftinds = leftinds[outsideinds]
                                localpoints = localpoints[outsideinds]
                            if len(leftinds) == 0:
                                break
                        continue
                    elif geom.GetType() == KinBody.Link.GeomType.Box:
                        insideinds = numpy.all(numpy.less_equal(abs(localpoints), geom.GetBoxExtents()) ,1)
                    elif geom.GetType() == KinBody.Link.GeomType.Sphere:
                        insideinds = numpy.less_equal(sum(localpoints**2,1), geom.GetSphereRadius()**2)
                    elif geom.GetType() == KinBody.Link.GeomType.Cylinder:
                        insideinds = numpy.less_equal(abs(localpoints[:,1]), 0.5*geom.GetCylinderHeight()) and numpy.less_equal(localpoint[:,0]**2+localpoint[:2]**2, geom.GetCylinderRadius()**2)
                    else:
                        continue
                    inside[leftinds[flatnonzero(insideinds)]] = True
                    outsideinds = flatnonzero(insideinds==0)
                    if len(outsideinds) > 0:
                        leftpoints = leftpoints[outsideinds]
                        leftinds = leftinds[outsideinds]
                    if len(leftinds) == 0:
                        break
        return inside
//  [python]
def test_all(type, shape):
    size = 1
    for i in range(len(shape)):
        size *= shape[i]

    for i in range(2**size):
        t = i

        a = numpy.empty(size, dtype=type)

        for j in range(size):
            a[j] = 0 if t % 2 == 0 else j + 1
            t = t >> 1

        a = a.reshape(shape)

        ia = dpnp.array(a)

        np_res = numpy.all(a)
        dpnp_res = dpnp.all(ia)
        assert_allclose(dpnp_res, np_res)

        np_res = a.all()
        dpnp_res = ia.all()
        assert_allclose(dpnp_res, np_res)
//  [python]
def _get_value(self, position) -> Optional[int]:
        """Index of the point at a given 2D position in data coordinates.

        Parameters
        ----------
        position : tuple
            Position in data coordinates.

        Returns
        -------
        value : int or None
            Index of point that is at the current coordinate if any.
        """
        # Display points if there are any in this slice
        view_data = self._view_data
        selection = None
        if len(view_data) > 0:
            displayed_position = [
                position[i] for i in self._slice_input.displayed
            ]
            # positions are scaled anisotropically by scale, but sizes are not,
            # so we need to calculate the ratio to correctly map to screen coordinates
            scale_ratio = (
                self.scale[self._slice_input.displayed] / self.scale[-1]
            )
            # Get the point sizes
            # TODO: calculate distance in canvas space to account for canvas_size_limits.
            # Without this implementation, point hover and selection (and anything depending
            # on self.get_value()) won't be aware of the real extent of points, causing
            # unexpected behaviour. See #3734 for details.
            sizes = np.expand_dims(self._view_size, axis=1) / scale_ratio / 2
            distances = abs(view_data - displayed_position)
            in_slice_matches = np.alltrue(
                distances <= sizes,
                axis=1,
            )
            indices = np.where(in_slice_matches)[0]
            if len(indices) > 0:
                selection = self._indices_view[indices[-1]]

        return selection
//  [python]
def _get_value(self, position) -> Optional[int]:
        """Index of the point at a given 2D position in data coordinates.

        Parameters
        ----------
        position : tuple
            Position in data coordinates.

        Returns
        -------
        value : int or None
            Index of point that is at the current coordinate if any.
        """
        # Display points if there are any in this slice
        view_data = self._view_data
        selection = None
        if len(view_data) > 0:
            displayed_position = [
                position[i] for i in self._slice_input.displayed
            ]
            # positions are scaled anisotropically by scale, but sizes are not,
            # so we need to calculate the ratio to correctly map to screen coordinates
            scale_ratio = (
                self.scale[self._slice_input.displayed] / self.scale[-1]
            )
            # Get the point sizes
            # TODO: calculate distance in canvas space to account for canvas_size_limits.
            # Without this implementation, point hover and selection (and anything depending
            # on self.get_value()) won't be aware of the real extent of points, causing
            # unexpected behaviour. See #3734 for details.
            sizes = np.expand_dims(self._view_size, axis=1) / scale_ratio / 2
            distances = abs(view_data - displayed_position)
            in_slice_matches = np.all(
                distances <= sizes,
                axis=1,
            )
            indices = np.where(in_slice_matches)[0]
            if len(indices) > 0:
                selection = self._indices_view[indices[-1]]

        return selection
//  [python]
def solve_theta(self, theta_point: numpy.ndarray) -> Optional[SolverOutput]:
        r"""
        Substitutes theta into the multiparametric problem and solves the following optimization problem

        .. math::

            \min_{x} \tilde{c}^Tx

        .. math::
            \begin{align}
            Ax &\leq \tilde{b}\\
            A_{eq}x &= \tilde{b}_{eq}\\
            x &\in R^n\\
            \end{align}

        :param theta_point: An uncertainty realization
        :return: The Solver output of the substituted problem, returns None if not solvable
        """

        if not n.all(self.A_t @ theta_point <= self.b_t):
            return None

        sol_obj = self.solver.solve_lp(c=self.H @ theta_point + self.c, A=self.A, b=self.b + self.F @ theta_point,
                                       equality_constraints=self.equality_indices)

        if sol_obj is not None:
            sol_obj.obj += self.c_c + self.c_t.T @ theta_point + 0.5 * theta_point.T @ self.Q_t @ theta_point
            return sol_obj

        return None
//  [python]
def solve_theta(self, theta_point: numpy.ndarray) -> Optional[SolverOutput]:
        r"""
        Substitutes theta into the multiparametric problem and solves the following optimization problem

        .. math::

            \min_{x} \tilde{c}^Tx

        .. math::
            \begin{align}
            Ax &\leq \tilde{b}\\
            A_{eq}x &= \tilde{b}_{eq}\\
            x &\in R^n\\
            \end{align}

        :param theta_point: An uncertainty realization
        :return: The Solver output of the substituted problem, returns None if not solvable
        """

        if not numpy.all(self.A_t @ theta_point <= self.b_t):
            return None

        sol_obj = self.solver.solve_lp(c=self.H @ theta_point + self.c, A=self.A, b=self.b + self.F @ theta_point,
                                       equality_constraints=self.equality_indices)

        if sol_obj is not None:
            sol_obj.obj += self.c_c + self.c_t.T @ theta_point + 0.5 * theta_point.T @ self.Q_t @ theta_point
            return sol_obj

        return None
//  [python]
def __init__(self, endog, exog, exog_coint, k_ar,
                 coint_rank, alpha, beta, gamma, sigma_u, deterministic='nc',
                 seasons=0, first_season=0, delta_y_1_T=None, y_lag1=None,
                 delta_x=None, model=None, names=None, dates=None):
        self.model = model
        self.y_all = endog
        self.exog = exog
        self.exog_coint = exog_coint
        self.names = names
        self.dates = dates
        self.neqs = endog.shape[0]
        self.k_ar = k_ar
        self.deterministic = deterministic
        self.seasons = seasons
        self.first_season = first_season

        self.coint_rank = coint_rank
        if alpha.dtype == np.complex128 and np.all(np.imag(alpha) == 0):
            alpha = np.real_if_close(alpha)
        if beta.dtype == np.complex128 and np.all(np.imag(beta) == 0):
            beta = np.real_if_close(beta)
        if gamma.dtype == np.complex128 and np.all(np.imag(gamma) == 0):
            gamma = np.real_if_close(gamma)

        self.alpha = alpha
        self.beta, self.det_coef_coint = np.vsplit(beta, [self.neqs])
        self.gamma, self.det_coef = np.hsplit(gamma,
                                              [self.neqs * (self.k_ar - 1)])

        if "ci" in deterministic:
            self.const_coint = self.det_coef_coint[:1, :]
        else:
            self.const_coint = np.zeros(coint_rank).reshape((1, -1))
        if "li" in deterministic:
            start = 1 if "ci" in deterministic else 0
            self.lin_trend_coint = self.det_coef_coint[start:start+1, :]
        else:
            self.lin_trend_coint = np.zeros(coint_rank).reshape(1, -1)
        if self.exog_coint is not None:
            start = ("ci" in deterministic) + ("li" in deterministic)
            self.exog_coint_coefs = self.det_coef_coint[start:, :]
        else:
            self.exog_coint_coefs = None

        split_const_season = 1 if "co" in deterministic else 0
        split_season_lin = split_const_season + ((seasons-1) if seasons else 0)
        if "lo" in deterministic:
            split_lin_exog = split_season_lin + 1
        else:
            split_lin_exog = split_season_lin
        self.const, self.seasonal, self.lin_trend, self.exog_coefs = \
            np.hsplit(self.det_coef,
                      [split_const_season, split_season_lin, split_lin_exog])

        self.sigma_u = sigma_u

        if y_lag1 is not None and delta_x is not None \
                and delta_y_1_T is not None:
            self._delta_y_1_T = delta_y_1_T
            self._y_lag1 = y_lag1
            self._delta_x = delta_x
        else:
            _y_1_T, self._delta_y_1_T, self._y_lag1, self._delta_x = \
                _endog_matrices(endog, self.exog, k_ar,
                                deterministic, seasons)
        self.nobs = self._y_lag1.shape[1]
//  [python]
def get_semantic_annotations(self, label, num_class_obj):
        annotation_classes = label["classes"]
        annotation_masks = label["masks"]

        texts = ["a semantic photo"] * self.num_text
        classes = []
        masks = []

        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx]
            if not n.all(mask is False):
                if class_id not in classes:
                    cls_name = self.metadata[str(class_id)]
                    classes.append(class_id)
                    masks.append(mask)
                    num_class_obj[cls_name] += 1
                else:
                    idx = classes.index(class_id)
                    masks[idx] += mask
                    masks[idx] = np.clip(masks[idx], 0, 1)

        num = 0
        for i, cls_name in enumerate(self.metadata["class_names"]):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f"a photo with a {cls_name}"
                    num += 1

        classes = np.array(classes)
        masks = np.array(masks)
        return classes, masks, texts
//  [python]
def get_semantic_annotations(self, label, num_class_obj):
        annotation_classes = label["classes"]
        annotation_masks = label["masks"]

        texts = ["a semantic photo"] * self.num_text
        classes = []
        masks = []

        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx]
            if not np.all(mask is False):
                if class_id not in classes:
                    cls_name = self.metadata[str(class_id)]
                    classes.append(class_id)
                    masks.append(mask)
                    num_class_obj[cls_name] += 1
                else:
                    idx = classes.index(class_id)
                    masks[idx] += mask
                    masks[idx] = np.clip(masks[idx], 0, 1)

        num = 0
        for i, cls_name in enumerate(self.metadata["class_names"]):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f"a photo with a {cls_name}"
                    num += 1

        classes = np.array(classes)
        masks = np.array(masks)
        return classes, masks, texts
//  [python]
def get_instance_annotations(self, label, num_class_obj):
        annotation_classes = label["classes"]
        annotation_masks = label["masks"]

        texts = ["an instance photo"] * self.num_text
        classes = []
        masks = []

        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx]

            if class_id in self.metadata["thing_ids"]:
                if not n.all(mask is False):
                    cls_name = self.metadata[str(class_id)]
                    classes.append(class_id)
                    masks.append(mask)
                    num_class_obj[cls_name] += 1

        num = 0
        for i, cls_name in enumerate(self.metadata["class_names"]):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f"a photo with a {cls_name}"
                    num += 1

        classes = np.array(classes)
        masks = np.array(masks)
        return classes, masks, texts
//  [python]
def get_instance_annotations(self, label, num_class_obj):
        annotation_classes = label["classes"]
        annotation_masks = label["masks"]

        texts = ["an instance photo"] * self.num_text
        classes = []
        masks = []

        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx]

            if class_id in self.metadata["thing_ids"]:
                if not np.all(mask is False):
                    cls_name = self.metadata[str(class_id)]
                    classes.append(class_id)
                    masks.append(mask)
                    num_class_obj[cls_name] += 1

        num = 0
        for i, cls_name in enumerate(self.metadata["class_names"]):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f"a photo with a {cls_name}"
                    num += 1

        classes = np.array(classes)
        masks = np.array(masks)
        return classes, masks, texts
//  [python]
def get_panoptic_annotations(self, label, num_class_obj):
        annotation_classes = label["classes"]
        annotation_masks = label["masks"]

        texts = ["an panoptic photo"] * self.num_text
        classes = []
        masks = []

        for idx in range(len(annotation_classes)):
            class_id = annotation_classes[idx]
            mask = annotation_masks[idx].data
            if not np.all(mask is False):
                cls_name = self.metadata[str(class_id)]
                classes.append(class_id)
                masks.append(mask)
                num_class_obj[cls_name] += 1

        num = 0
        for i, cls_name in enumerate(self.metadata["class_names"]):
            if num_class_obj[cls_name] > 0:
                for _ in range(num_class_obj[cls_name]):
                    if num >= len(texts):
                        break
                    texts[num] = f"a photo with a {cls_name}"
                    num += 1

        classes = np.array(classes)
        masks = np.array(masks)
        return classes, masks, texts
//  [python]
def _lazyselect(condlist, choicelist, arrays, default=0):
    """
    Mimic `np.select(condlist, choicelist)`.

    Notice, it assumes that all `arrays` are of the same shape or can be
    broadcasted together.

    All functions in `choicelist` must accept array arguments in the order
    given in `arrays` and must return an array of the same shape as broadcasted
    `arrays`.

    Examples
    --------
    >>> import numpy as np
    >>> x = np.arange(6)
    >>> np.select([x <3, x > 3], [x**2, x**3], default=0)
    array([  0,   1,   4,   0,  64, 125])

    >>> _lazyselect([x < 3, x > 3], [lambda x: x**2, lambda x: x**3], (x,))
    array([   0.,    1.,    4.,   0.,   64.,  125.])

    >>> a = -np.ones_like(x)
    >>> _lazyselect([x < 3, x > 3],
    ...             [lambda x, a: x**2, lambda x, a: a * x**3],
    ...             (x, a), default=np.nan)
    array([   0.,    1.,    4.,   nan,  -64., -125.])

    """
    arrays = np.broadcast_arrays(*arrays)
    tcode = np.mintypecode([a.dtype.char for a in arrays])
    out = np.full(np.shape(arrays[0]), fill_value=default, dtype=tcode)
    for func, cond in zip(choicelist, condlist):
        if np.all(cond is False):
            continue
        cond, _ = np.broadcast_arrays(cond, arrays[0])
        temp = tuple(np.extract(cond, arr) for arr in arrays)
        np.place(out, cond, func(*temp))
    return out
//  [python]
def nuisance_estimator_generator(propensity_model, model_Y, random_state=None, second_stage=False):
        """Generate nuissance estimator given model inputs from the class."""
        def nuisance_estimator(Y, T, X, W, sample_weight=None, split_indices=None):
            # Expand one-hot encoding to include the zero treatment
            ohe_T = np.hstack([np.all(1 - T, axis=1, keepdims=True), T])
            # Test that T contains all treatments. If not, return None
            T = ohe_T @ np.arange(ohe_T.shape[1])
            if len(np.unique(T)) < ohe_T.shape[1]:
                return None
            # Nuissance estimates evaluated with cross-fitting
            this_random_state = check_random_state(random_state)
            if (split_indices is None) and second_stage:
                # Define 2-fold iterator
                kfold_it = StratifiedKFold(n_splits=2, shuffle=True, random_state=this_random_state).split(X, T)
                # Check if there is only one example of some class
                with warnings.catch_warnings():
                    warnings.filterwarnings('error')
                    try:
                        split_indices = list(kfold_it)[0]
                    except Warning as warn:
                        msg = str(warn)
                        if "The least populated class in y has only 1 members" in msg:
                            return None
            if W is not None:
                X_tilde = np.concatenate((X, W), axis=1)
            else:
                X_tilde = X
            try:
                if not second_stage:
                    # No need to crossfit for internal nodes
                    propensity_model_clone = clone(propensity_model, safe=False)
                    propensity_model_clone.fit(X_tilde, T)
                    propensities = propensity_model_clone.predict_proba(X_tilde)
                    Y_hat = _group_predict(X_tilde, ohe_T.shape[1],
                                           clone(model_Y, safe=False).fit(np.hstack([X_tilde, ohe_T]), Y).predict)
                else:
                    propensities = _cross_fit(propensity_model, X_tilde, T, split_indices,
                                              sample_weight=sample_weight, predict_func_name='predict_proba')
                    Y_hat = _group_cross_fit(model_Y, X_tilde, Y, ohe_T, split_indices, sample_weight=sample_weight)
            except ValueError as exc:
                raise ValueError("The original error: {0}".format(str(exc)) +
                                 " This might be caused by too few sample in the tree leafs." +
                                 " Try increasing the min_leaf_size.")
            return Y_hat, propensities
        return nuisance_estimator
//  [python]
def test_single_for_numpy():
    """
    Feature: JIT Fallback
    Description: Test fallback with control flow.
    Expectation: No exception.
    """
    @jit
    def control_flow_for():
        x = np.array([1, 3, 5])
        y = np.array([0, 2, 4])
        for _ in range(3):
            x = x + y
        if n.all(x == np.array([1, 9, 17])):
            return Tensor(x)
        return Tensor(x - 1)
    res = control_flow_for()
    assert (res.asnumpy() == [1, 9, 17]).all()
//  [python]
def test_single_for_numpy():
    """
    Feature: JIT Fallback
    Description: Test fallback with control flow.
    Expectation: No exception.
    """
    @jit
    def control_flow_for():
        x = np.array([1, 3, 5])
        y = np.array([0, 2, 4])
        for _ in range(3):
            x = x + y
        if np.all(x == np.array([1, 9, 17])):
            return Tensor(x)
        return Tensor(x - 1)
    res = control_flow_for()
    assert (res.asnumpy() == [1, 9, 17]).all()
//  [python]
def test_intersect(lhs, rhs, mask, expected):
    if mask is None:
        mask = np.uint64(0xFFFFFFFFFFFFFFFF)
    lhs_idx, rhs_idx = intersect(lhs, rhs, mask)
    result = lhs[lhs_idx] & mask
    assert n.all(result == expected)
//  [python]
def test_intersect(lhs, rhs, mask, expected):
    if mask is None:
        mask = np.uint64(0xFFFFFFFFFFFFFFFF)
    lhs_idx, rhs_idx = intersect(lhs, rhs, mask)
    result = lhs[lhs_idx] & mask
    assert np.all(result == expected)
//  [python]
def test_intersect_keep_both(lhs, rhs, mask, expected):
    if mask is None:
        mask = np.uint64(0xFFFFFFFFFFFFFFFF)
    lhs_idx, rhs_idx = intersect(lhs, rhs, mask, drop_duplicates=False)
    expected_lhs = np.argwhere(np.in1d(lhs, rhs)).flatten()
    expected_rhs = np.argwhere(np.in1d(rhs, lhs)).flatten()
    assert n.all(lhs_idx == expected_lhs)
    assert n.all(rhs_idx == expected_rhs)
//  [python]
def test_intersect_keep_both(lhs, rhs, mask, expected):
    if mask is None:
        mask = np.uint64(0xFFFFFFFFFFFFFFFF)
    lhs_idx, rhs_idx = intersect(lhs, rhs, mask, drop_duplicates=False)
    expected_lhs = np.argwhere(np.in1d(lhs, rhs)).flatten()
    expected_rhs = np.argwhere(np.in1d(rhs, lhs)).flatten()
    assert np.all(lhs_idx == expected_lhs)
    assert np.all(rhs_idx == expected_rhs)
//  [python]
def w_search_scenarios(lst):
    """Decorate with array and every possible target."""
    params = []
    for arr in lst:
        for idx in range(len(arr)):
            np_arr = u64(arr)
            # assert sorted
            assert n.all(np_arr[:-1] <= np_arr[1:])
            params.append((u64(arr),
                           u64(arr[idx]),
                           (u64(idx), True)))
        # Add one before and one after
        params.append((u64(arr), u64(arr[0] - 1), (u64(0), False)))
        params.append((u64(arr), u64(arr[-1] + 1), (u64(len(arr)), False)))
    return pytest.mark.parametrize(
        "array,target,expected", params
    )
//  [python]
def w_search_scenarios(lst):
    """Decorate with array and every possible target."""
    params = []
    for arr in lst:
        for idx in range(len(arr)):
            np_arr = u64(arr)
            # assert sorted
            assert np.all(np_arr[:-1] <= np_arr[1:])
            params.append((u64(arr),
                           u64(arr[idx]),
                           (u64(idx), True)))
        # Add one before and one after
        params.append((u64(arr), u64(arr[0] - 1), (u64(0), False)))
        params.append((u64(arr), u64(arr[-1] + 1), (u64(len(arr)), False)))
    return pytest.mark.parametrize(
        "array,target,expected", params
    )
//  [python]
def actualise_lambda(l, val, new_val, A, erreur, new_J, old_J):
    dim = len(val)
    id = NP.identity(dim)
    # Matrice du système
    Q = NP.dot(NP.transpose(A), A) + l * id
    # Second membre du système
    d = NP.dot(NP.transpose(A), erreur)
    old_Q = old_J
    new_Q = old_J + 0.5 * \
        NP.dot(NP.transpose(new_val - val), NP.dot(Q, new_val - val)) + \
        NP.dot(NP.transpose(new_val - val), d)
    # Ratio de la décroissance réelle et de l'approx. quad.
    if n.all((old_Q - new_Q) != 0.):
        R = (old_J - new_J) / (old_Q - new_Q)
        if (R < 0.25):
            l = l * 10.
        elif (R > 0.75):
            l = l / 15.
    else:
        if (old_J > new_J):
            l = l * 10.
        else:
            l = l / 10.
    return l
//  [python]
def actualise_lambda(l, val, new_val, A, erreur, new_J, old_J):
    dim = len(val)
    id = NP.identity(dim)
    # Matrice du système
    Q = NP.dot(NP.transpose(A), A) + l * id
    # Second membre du système
    d = NP.dot(NP.transpose(A), erreur)
    old_Q = old_J
    new_Q = old_J + 0.5 * \
        NP.dot(NP.transpose(new_val - val), NP.dot(Q, new_val - val)) + \
        NP.dot(NP.transpose(new_val - val), d)
    # Ratio de la décroissance réelle et de l'approx. quad.
    if NP.all((old_Q - new_Q) != 0.):
        R = (old_J - new_J) / (old_Q - new_Q)
        if (R < 0.25):
            l = l * 10.
        elif (R > 0.75):
            l = l / 15.
    else:
        if (old_J > new_J):
            l = l * 10.
        else:
            l = l / 10.
    return l
//  [python]
def test_unique(array):
    expected = np.unique(array)
    result = unique(array)
    assert np.all(result == expected)
//  [python]
def intersect_curves(nodes1, nodes2):
    r"""Intersect two parametric B |eacute| zier curves.

    Args:
        nodes1 (numpy.ndarray): The nodes in the first curve.
        nodes2 (numpy.ndarray): The nodes in the second curve.

    Returns:
        numpy.ndarray: ``2 x N`` array of intersection parameters.
        Each row contains a pair of values :math:`s` and :math:`t`
        (each in :math:`\left[0, 1\right]`) such that the curves
        intersect: :math:`B_1(s) = B_2(t)`.

    Raises:
        NotImplementedError: If the "intersection polynomial" is
            all zeros -- which indicates coincident curves.
    """
    nodes1 = _curve_helpers.full_reduce(nodes1)
    nodes2 = _curve_helpers.full_reduce(nodes2)
    _, num_nodes1 = nodes1.shape
    _, num_nodes2 = nodes2.shape
    swapped = False
    if num_nodes1 > num_nodes2:
        nodes1, nodes2 = nodes2, nodes1
        swapped = True
    coeffs = normalize_polynomial(to_power_basis(nodes1, nodes2))
    if n.all(coeffs == 0.0):
        raise NotImplementedError(_COINCIDENT_ERR)

    _check_non_simple(coeffs)
    t_vals = roots_in_unit_interval(coeffs)
    final_s = []
    final_t = []
    for t_val in t_vals:
        (x_val,), (y_val,) = _curve_helpers.evaluate_multi(
            nodes2, np.asfortranarray([t_val])
        )
        s_val = locate_point(nodes1, x_val, y_val)
        if s_val is not None:
            _resolve_and_add(nodes1, s_val, final_s, nodes2, t_val, final_t)
    result = np.zeros((2, len(final_s)), order="F")
    if swapped:
        final_s, final_t = final_t, final_s
    result[0, :] = final_s
    result[1, :] = final_t
    return result
//  [python]
def intersect_curves(nodes1, nodes2):
    r"""Intersect two parametric B |eacute| zier curves.

    Args:
        nodes1 (numpy.ndarray): The nodes in the first curve.
        nodes2 (numpy.ndarray): The nodes in the second curve.

    Returns:
        numpy.ndarray: ``2 x N`` array of intersection parameters.
        Each row contains a pair of values :math:`s` and :math:`t`
        (each in :math:`\left[0, 1\right]`) such that the curves
        intersect: :math:`B_1(s) = B_2(t)`.

    Raises:
        NotImplementedError: If the "intersection polynomial" is
            all zeros -- which indicates coincident curves.
    """
    nodes1 = _curve_helpers.full_reduce(nodes1)
    nodes2 = _curve_helpers.full_reduce(nodes2)
    _, num_nodes1 = nodes1.shape
    _, num_nodes2 = nodes2.shape
    swapped = False
    if num_nodes1 > num_nodes2:
        nodes1, nodes2 = nodes2, nodes1
        swapped = True
    coeffs = normalize_polynomial(to_power_basis(nodes1, nodes2))
    if np.all(coeffs == 0.0):
        raise NotImplementedError(_COINCIDENT_ERR)

    _check_non_simple(coeffs)
    t_vals = roots_in_unit_interval(coeffs)
    final_s = []
    final_t = []
    for t_val in t_vals:
        (x_val,), (y_val,) = _curve_helpers.evaluate_multi(
            nodes2, np.asfortranarray([t_val])
        )
        s_val = locate_point(nodes1, x_val, y_val)
        if s_val is not None:
            _resolve_and_add(nodes1, s_val, final_s, nodes2, t_val, final_t)
    result = np.zeros((2, len(final_s)), order="F")
    if swapped:
        final_s, final_t = final_t, final_s
    result[0, :] = final_s
    result[1, :] = final_t
    return result
//  [python]
def test_unique_shifted(array, shift):
    expected = np.unique(array >> shift)
    result = unique(array, shift)
    assert np.all(result == expected)
//  [python]
def test_same_as_numpy(seed):
    np.random.seed(seed)
    rand_arr_1 = np.random.randint(0, 500, 100, dtype=np.uint64)
    rand_arr_2 = np.random.randint(0, 500, 100, dtype=np.uint64)
    rand_arr_1.sort()
    rand_arr_2.sort()

    expected_result = np.intersect1d(rand_arr_1, rand_arr_2)

    lhs_idx, rhs_idx = intersect(rand_arr_1, rand_arr_2)
    result = rand_arr_1[lhs_idx]
    assert np.alltrue(result == expected_result)
//  [python]
def test_same_as_numpy(seed):
    np.random.seed(seed)
    rand_arr_1 = np.random.randint(0, 500, 100, dtype=np.uint64)
    rand_arr_2 = np.random.randint(0, 500, 100, dtype=np.uint64)
    rand_arr_1.sort()
    rand_arr_2.sort()

    expected_result = np.intersect1d(rand_arr_1, rand_arr_2)

    lhs_idx, rhs_idx = intersect(rand_arr_1, rand_arr_2)
    result = rand_arr_1[lhs_idx]
    assert np.all(result == expected_result)
//  [python]
def test_merge(lhs, rhs, expected):
    result = merge(lhs, rhs)
    assert n.all(result == expected)
//  [python]
def test_merge(lhs, rhs, expected):
    result = merge(lhs, rhs)
    assert np.all(result == expected)
//  [python]
def test_unique_matches_snp(seed):
    np.random.seed(seed)
    rand_arr_1 = np.random.randint(0, 500, 100, dtype=np.uint64)
    rand_arr_1.sort()
    expected = np.unique(rand_arr_1)
    result = unique(rand_arr_1)
    assert (result == expected).all()
//  [python]
def test_unique_matches_snp(seed):
    np.random.seed(seed)
    rand_arr_1 = np.random.randint(0, 500, 100, dtype=np.uint64)
    rand_arr_1.sort()
    expected = np.unique(rand_arr_1)
    result = unique(rand_arr_1)
    assert np.all(result == expected)
//  [python]
def test_merge_w_drop(lhs, rhs, expected):
    result = merge(lhs, rhs, drop_duplicates=True)
    assert np.all(result == np.sort(np.unique(expected)))
//  [python]
def test_adjacent(lhs, rhs, mask, delta, lhs_idx_expected, rhs_idx_expected):
    if mask is None:
        mask = np.uint64(0xFFFFFFFFFFFFFFFF)
    if delta == -1:
        rhs_idx, lhs_idx = adjacent(rhs, lhs, mask)
    else:
        lhs_idx, rhs_idx = adjacent(lhs, rhs, mask)
    assert n.all(lhs_idx == lhs_idx_expected)
    assert n.all(rhs_idx == rhs_idx_expected)