jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_bare_trainonlytransformer_pandas(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used directly"""
    X, y = (
        pd.DataFrame(random_xy_dataset_regr[0]),
        pd.DataFrame(random_xy_dataset_regr[1]),
    )
    X_train, X_test, y_train, y_test = train_test_split(X, y)

    trf = TrainOnlyTrainOnlyTransformer()
    trf.fit(X_train, y_train)

    assert np.all(trf.transform(X_train) != X_train)
    assert np.all(trf.transform(X_test) == X_test)

    assert isinstance(trf.transform(X_train), pd.DataFrame)
    assert isinstance(trf.transform(X_test), pd.DataFrame)
//  [python]
def test_pipeline_trainonlytransformer_pandas(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used in a pipeline"""
    X, y = (
        pd.DataFrame(random_xy_dataset_regr[0]),
        pd.DataFrame(random_xy_dataset_regr[1]),
    )
    X_train, X_test, y_train, y_test = train_test_split(X, y)

    trf = make_pipeline(TrainOnlyTrainOnlyTransformer())
    trf.fit(X_train, y_train)

    assert np.all(trf.transform(X_train) != X_train)
    assert np.all(trf.transform(X_test) == X_test)

    assert isinstance(trf.transform(X_train), pd.DataFrame)
    assert isinstance(trf.transform(X_test), pd.DataFrame)
//  [python]
def test_ufunc_at_negative(self):
        arr = np.ones(5, dtype=np.int32)
        indx = np.arange(5)
        umt.indexed_negative.at(arr, indx)
        # If it is [-1, -1, -1, -100, 0] then the regular strided loop was used
        assert n.all(arr == [-1, -1, -1, -200, -1])
//  [python]
def test_ufunc_at_negative(self):
        arr = np.ones(5, dtype=np.int32)
        indx = np.arange(5)
        umt.indexed_negative.at(arr, indx)
        # If it is [-1, -1, -1, -100, 0] then the regular strided loop was used
        assert np.all(arr == [-1, -1, -1, -200, -1])
//  [python]
def test_sign_magnitude_positive_definite(self):
    for dtype in [dtypes.float32, dtypes.float64]:
      with self.cached_session():
        matrix_tensor = math_utils.sign_magnitude_positive_definite(
            raw=constant_op.constant([[-1., -2.], [3., 4.]], dtype=dtype),
            off_diagonal_scale=constant_op.constant(-1., dtype=dtype),
            overall_scale=constant_op.constant(1., dtype=dtype))
        matrix_evaled = matrix_tensor.eval()
        self.assertAllClose(matrix_evaled, matrix_evaled.T)
        self.assertTrue(numpy.all(numpy.linalg.eigvals(matrix_evaled) > 0))
//  [python]
def transform(self, point):
        """Match a `point` containing integers to real vector representations of them.

        If the upper bound of integers supported by an instance of `OneHotEncode`
        is less or equal to 2, then cast them to floats.

        .. note:: This transformation possibly appends one more tensor dimension to `point`.
        """
        point_ = numpy.asarray(point)
        assert (
            n.all(point_ < self.num_cats)
            and n.all(point_ >= 0)
            and n.all(point_ % 1 == 0)
        )

        if self.num_cats <= 2:
            return numpy.asarray(point_, dtype=float)

        hot = numpy.zeros(self.infer_target_shape(point_.shape))
        grid = numpy.meshgrid(
            *[numpy.arange(dim) for dim in point_.shape], indexing="ij"
        )
        hot[tuple(grid + [point_])] = 1
        return hot
//  [python]
def transform(self, point):
        """Match a `point` containing integers to real vector representations of them.

        If the upper bound of integers supported by an instance of `OneHotEncode`
        is less or equal to 2, then cast them to floats.

        .. note:: This transformation possibly appends one more tensor dimension to `point`.
        """
        point_ = numpy.asarray(point)
        assert (
            numpy.all(point_ < self.num_cats)
            and numpy.all(point_ >= 0)
            and numpy.all(point_ % 1 == 0)
        )

        if self.num_cats <= 2:
            return numpy.asarray(point_, dtype=float)

        hot = numpy.zeros(self.infer_target_shape(point_.shape))
        grid = numpy.meshgrid(
            *[numpy.arange(dim) for dim in point_.shape], indexing="ij"
        )
        hot[tuple(grid + [point_])] = 1
        return hot
//  [python]
def test_at_negative_indexes(self, dtype, ufunc):
        a = np.arange(0, 10).astype(dtype)
        indxs = np.array([-1, 1, -1, 2]).astype(np.intp)
        vals = np.array([1, 5, 2, 10], dtype=a.dtype)

        expected = a.copy()
        for i, v in zip(indxs, vals):
            expected[i] = ufunc(expected[i], v)

        ufunc.at(a, indxs, vals)
        assert_array_equal(a, expected)
        assert np.all(indxs == [-1, 1, -1, 2])
//  [python]
def test_find_access_past_buffer():
    # This checks that no read past the string buffer occurs in
    # string_fastsearch.h. The buffer class makes sure this is checked.
    # To see it in action, you can remove the checks in the buffer and
    # this test will produce an 'Invalid read' if run under valgrind.
    arr = np.array([b'abcd', b'ebcd'])
    result = np._core.umath.find(arr, b'cde', 0, np.iinfo(np.int64).max)
    assert np.all(result == -1)
//  [python]
def test_complexvalued(tmp_path):
    freq = unyt_array([1j, 1j * 10], "Hz")
    arr = 1 / (Unit("F") * Unit("Ω") * freq)
    arr = arr.to("dimensionless")
    assert arr.units.is_dimensionless
    assert np.all(arr.v == np.asarray([-1j, -1j * 0.1]))
    arr = unyt_array([1j, 1j * 10], "mJ")
    arr.convert_to_base()
    assert_allclose_units(arr, unyt_array([1j * 0.001, 1j * 0.01], "J"))
    arr.convert_to_units("mJ")
    assert_allclose_units(arr, unyt_array([1j, 1j * 10], "mJ"))
    arr.convert_to_mks()
    assert_allclose_units(arr, unyt_array([1j * 0.001, 1j * 0.01], "J"))
    arr.convert_to_cgs()
    assert_allclose_units(arr, unyt_array([1j * 10000, 1j * 100000], "erg"))
    arr.convert_to_equivalent("K", "thermal")
    assert_allclose_units(
        arr, unyt_array([1j * 7.24297157e19, 1j * 7.24297157e20], "K")
    )
    arr = arr.to_equivalent("J", "thermal")
    assert_allclose_units(arr, unyt_array([1j * 0.001, 1j * 0.01], "J"))
    assert_allclose_units(arr.to_ndarray(), np.asarray([1j * 0.001, 1j * 0.01]))
    assert_allclose_units(arr.to_value(), np.asarray([1j * 0.001, 1j * 0.01]))
    assert arr.tolist() == [1j * 0.001, 1j * 0.01]
    assert_allclose_units(arr.in_units("mJ"), unyt_array([1j, 1j * 10], "mJ"))
    assert_allclose_units(arr.in_base(), unyt_array([1j * 0.001, 1j * 0.01], "J"))
    assert_allclose_units(arr.in_cgs(), unyt_array([1j * 10000, 1j * 100000], "erg"))
    assert_allclose_units(arr.in_mks(), unyt_array([1j * 0.001, 1j * 0.01], "J"))
    fname = tmp_path / "testcomplexvalued.txt"
    savetxt(fname, arr)
    farr = loadtxt(fname, dtype=np.complex128)
    assert_allclose_units(farr, unyt_array([1j * 0.001, 1j * 0.01], "J"))
//  [python]
def test_random_sparse_signed_matrix():
    # Define parameter
    p = .33
    p_e = .75
    mean_E = .4
    std_E = 0
    n_in = 400
    neuron_sign = rd.choice([1, -1], n_in, p=[p_e, 1 - p_e])

    M1, M1_sign, M1_theta, M1_is_con = random_sparse_signed_matrix(neuron_sign=neuron_sign, p=p,
                                                                   balance_zero_mean_per_neuron=True)
    s1, _ = la.eig(M1)

    assert n.all(n.abs(M1[M1_is_con]) == M1_theta[M1_is_con])
    assert n.all(n.sign(M1) == M1_sign * M1_is_con)
    assert n.all(M1_is_con == (M1_theta > 0))

    M2, _, _, _ = random_sparse_signed_matrix(neuron_sign=neuron_sign, p=1., balance_zero_mean_per_neuron=True)
    M2 = M2 * (rd.rand(n_in, n_in) < p)
    s2, _ = la.eig(M2)

    fig, ax_list = plt.subplots(2)

    ax_list[0].set_title('Random sign constrained without neuron specific balance (p={:.3g})'.format(p))
    ax_list[1].set_title('Random sign constrained, probability mask taken after scaling')
    ax_list[0].scatter(s1.real, s1.imag)
    ax_list[1].scatter(s2.real, s2.imag)
    c = plt.Circle(xy=(0, 0), radius=1, edgecolor='r', alpha=.5)
    ax_list[0].add_artist(c)
    c = plt.Circle(xy=(0, 0), radius=1, edgecolor='r', alpha=.5)
    ax_list[1].add_artist(c)

    for ax in ax_list:
        ax.set_xlim([-2, 2])
        ax.set_ylim([-2, 2])

    plt.show()
//  [python]
def test_random_sparse_signed_matrix():
    # Define parameter
    p = .33
    p_e = .75
    mean_E = .4
    std_E = 0
    n_in = 400
    neuron_sign = rd.choice([1, -1], n_in, p=[p_e, 1 - p_e])

    M1, M1_sign, M1_theta, M1_is_con = random_sparse_signed_matrix(neuron_sign=neuron_sign, p=p,
                                                                   balance_zero_mean_per_neuron=True)
    s1, _ = la.eig(M1)

    assert np.all(np.abs(M1[M1_is_con]) == M1_theta[M1_is_con])
    assert np.all(np.sign(M1) == M1_sign * M1_is_con)
    assert np.all(M1_is_con == (M1_theta > 0))

    M2, _, _, _ = random_sparse_signed_matrix(neuron_sign=neuron_sign, p=1., balance_zero_mean_per_neuron=True)
    M2 = M2 * (rd.rand(n_in, n_in) < p)
    s2, _ = la.eig(M2)

    fig, ax_list = plt.subplots(2)

    ax_list[0].set_title('Random sign constrained without neuron specific balance (p={:.3g})'.format(p))
    ax_list[1].set_title('Random sign constrained, probability mask taken after scaling')
    ax_list[0].scatter(s1.real, s1.imag)
    ax_list[1].scatter(s2.real, s2.imag)
    c = plt.Circle(xy=(0, 0), radius=1, edgecolor='r', alpha=.5)
    ax_list[0].add_artist(c)
    c = plt.Circle(xy=(0, 0), radius=1, edgecolor='r', alpha=.5)
    ax_list[1].add_artist(c)

    for ax in ax_list:
        ax.set_xlim([-2, 2])
        ax.set_ylim([-2, 2])

    plt.show()
//  [python]
def pie(self, X, win=None, env=None, opts=None):
        """
        This function draws a pie chart based on the `N` tensor `X`.

        The following `opts` are supported:

        - `opts.legend`: `list` containing legend names
        """

        X = np.squeeze(X)
        assert X.ndim == 1, "X should be one-dimensional"
        assert n.all(np.greater_equal(X, 0)), "X cannot contain negative values"

        opts = {} if opts is None else opts
        _title2str(opts)
        _assert_opts(opts)

        data = [
            {
                "values": X.tolist(),
                "labels": opts.get("legend"),
                "type": "pie",
            }
        ]
        return self._send(
            {
                "data": data,
                "win": win,
                "eid": env,
                "layout": _opts2layout(opts),
                "opts": opts,
            }
        )
//  [python]
def pie(self, X, win=None, env=None, opts=None):
        """
        This function draws a pie chart based on the `N` tensor `X`.

        The following `opts` are supported:

        - `opts.legend`: `list` containing legend names
        """

        X = np.squeeze(X)
        assert X.ndim == 1, "X should be one-dimensional"
        assert np.all(np.greater_equal(X, 0)), "X cannot contain negative values"

        opts = {} if opts is None else opts
        _title2str(opts)
        _assert_opts(opts)

        data = [
            {
                "values": X.tolist(),
                "labels": opts.get("legend"),
                "type": "pie",
            }
        ]
        return self._send(
            {
                "data": data,
                "win": win,
                "eid": env,
                "layout": _opts2layout(opts),
                "opts": opts,
            }
        )
//  [python]
def __init__(self, endog, exog, **kwargs):
        super(BinaryModel, self).__init__(endog, exog, **kwargs)
        if (not issubclass(self.__class__, MultinomialModel) and
                not np.all((self.endog >= 0) & (self.endog <= 1))):
            raise ValueError("endog must be in the unit interval.")
//  [python]
def top_k(cls, harness):

    def custom_assert(tst, result_jax, result_tf, *, err_msg, **_):
      assert len(result_jax) == len(result_tf)
      # TODO: TF and JAX sort [inf, nan] differently.
      first_arr_jax, first_arr_tf = result_jax[0], result_tf[0]
      if n.all(first_arr_jax == first_arr_tf):
        for arr_jax, arr_tf in zip(result_jax, result_tf):
          tst.assertArraysEqual(arr_jax, arr_tf, err_msg=err_msg)
      else:
        mask_jax = np.isnan(first_arr_jax) | np.isinf(first_arr_jax)
        mask_tf = np.isnan(first_arr_tf) | np.isinf(first_arr_tf)
        tst.assertArraysEqual(
            first_arr_jax[~mask_jax], first_arr_tf[~mask_tf], err_msg=err_msg)

    return [
        custom_numeric(
            dtypes=[np.float16, dtypes.bfloat16, np.float32, np.float64],
            custom_assert=custom_assert,
            description=(
                "Produces different results when the array contains `inf` and `NaN`"
                " (they are sorted differently in TF vs. XLA)."))
    ]
//  [python]
def top_k(cls, harness):

    def custom_assert(tst, result_jax, result_tf, *, err_msg, **_):
      assert len(result_jax) == len(result_tf)
      # TODO: TF and JAX sort [inf, nan] differently.
      first_arr_jax, first_arr_tf = result_jax[0], result_tf[0]
      if np.all(first_arr_jax == first_arr_tf):
        for arr_jax, arr_tf in zip(result_jax, result_tf):
          tst.assertArraysEqual(arr_jax, arr_tf, err_msg=err_msg)
      else:
        mask_jax = np.isnan(first_arr_jax) | np.isinf(first_arr_jax)
        mask_tf = np.isnan(first_arr_tf) | np.isinf(first_arr_tf)
        tst.assertArraysEqual(
            first_arr_jax[~mask_jax], first_arr_tf[~mask_tf], err_msg=err_msg)

    return [
        custom_numeric(
            dtypes=[np.float16, dtypes.bfloat16, np.float32, np.float64],
            custom_assert=custom_assert,
            description=(
                "Produces different results when the array contains `inf` and `NaN`"
                " (they are sorted differently in TF vs. XLA)."))
    ]
//  [python]
def perform(self, node, inputs, out_):
        out, = out_
        v = inputs[0]
        out[0] = v
        assert n.all(inputs[1:]), self.msg
//  [python]
def perform(self, node, inputs, out_):
        out, = out_
        v = inputs[0]
        out[0] = v
        assert numpy.all(inputs[1:]), self.msg
//  [python]
def make_node(self, value, *conds):
        if not isinstance(value, Variable):
            value = T.as_tensor_variable(value)
        cond = [T.as_tensor_variable(c) for c in conds]
        assert n.all([c.type.ndim == 0 for c in cond])
        return gof.Apply(self, [value] + cond, [value.type()])
//  [python]
def make_node(self, value, *conds):
        if not isinstance(value, Variable):
            value = T.as_tensor_variable(value)
        cond = [T.as_tensor_variable(c) for c in conds]
        assert numpy.all([c.type.ndim == 0 for c in cond])
        return gof.Apply(self, [value] + cond, [value.type()])
//  [python]
def __init__(self, param_grid, n_evaluations=10, random_state=None, start_evaluations=3,
                 subgrid_size=3, maximize=True):
        AbstractParameterGenerator.__init__(self, param_grid=param_grid, n_evaluations=n_evaluations,
                                            random_state=random_state, maximize=maximize)
        self.start_evaluations = start_evaluations
        self.subgrid_size = subgrid_size
        self.dimensions_sum = sum(self.dimensions)
        self.subgrid_parameter_generator = None
        if not n.all(n.array(self.dimensions) <= 2 * self.subgrid_size):
            logger = logging.getLogger(__name__)
            logger.info("Optimizing on subgrid")
            param_subgrid, self.subgrid_indices = _create_subgrid(self.param_grid, self.subgrid_size)
            self.subgrid_parameter_generator = \
                SubgridParameterOptimizer(param_subgrid, n_evaluations=self.n_evaluations // 2,
                                          subgrid_size=subgrid_size)
//  [python]
def __init__(self, param_grid, n_evaluations=10, random_state=None, start_evaluations=3,
                 subgrid_size=3, maximize=True):
        AbstractParameterGenerator.__init__(self, param_grid=param_grid, n_evaluations=n_evaluations,
                                            random_state=random_state, maximize=maximize)
        self.start_evaluations = start_evaluations
        self.subgrid_size = subgrid_size
        self.dimensions_sum = sum(self.dimensions)
        self.subgrid_parameter_generator = None
        if not numpy.all(numpy.array(self.dimensions) <= 2 * self.subgrid_size):
            logger = logging.getLogger(__name__)
            logger.info("Optimizing on subgrid")
            param_subgrid, self.subgrid_indices = _create_subgrid(self.param_grid, self.subgrid_size)
            self.subgrid_parameter_generator = \
                SubgridParameterOptimizer(param_subgrid, n_evaluations=self.n_evaluations // 2,
                                          subgrid_size=subgrid_size)
//  [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 numpy.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(numpy.multiply, num, one) / reduce(numpy.multiply, denum, one)
    if aslist:
        if numpy.alltrue(v == 1):
            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 numpy.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(numpy.multiply, num, one) / reduce(numpy.multiply, denum, one)
    if aslist:
        if numpy.all(v == 1):
            return []
        else:
            return [v]
    return v
//  [python]
def local_useless_rebroadcast(node):
    """
    Remove Rebroadcast if id does not actually change the broadcasting pattern.

    """
    if isinstance(node.op, T.Rebroadcast):
        x = node.inputs[0]
        if n.all(x.broadcastable == node.outputs[0].broadcastable):
            # No broadcastable flag was modified
            # No need to copy over stack trace,
            # because x should already have a stack trace.
            return [x]
        else:
            # Keep the flags that modify something
            new_axis = {}
            for dim, bc in list(node.op.axis.items()):
                if x.broadcastable[dim] != bc:
                    new_axis[dim] = bc
            if new_axis == node.op.axis:
                # All flags are useful
                return
            else:
                r = T.Rebroadcast(*list(new_axis.items()))(x)
                # Copy over stacktrace from previous output
                copy_stack_trace(node.outputs, r)
                return [r]
//  [python]
def local_useless_rebroadcast(node):
    """
    Remove Rebroadcast if id does not actually change the broadcasting pattern.

    """
    if isinstance(node.op, T.Rebroadcast):
        x = node.inputs[0]
        if numpy.all(x.broadcastable == node.outputs[0].broadcastable):
            # No broadcastable flag was modified
            # No need to copy over stack trace,
            # because x should already have a stack trace.
            return [x]
        else:
            # Keep the flags that modify something
            new_axis = {}
            for dim, bc in list(node.op.axis.items()):
                if x.broadcastable[dim] != bc:
                    new_axis[dim] = bc
            if new_axis == node.op.axis:
                # All flags are useful
                return
            else:
                r = T.Rebroadcast(*list(new_axis.items()))(x)
                # Copy over stacktrace from previous output
                copy_stack_trace(node.outputs, r)
                return [r]
//  [python]
def _maybe_reorder_labels(
        self,
        intermediate: "PandasDataframe",
        row_positions,
        sorted_row_positions,
        col_positions,
        sorted_col_positions,
    ) -> "PandasDataframe":
        """
        Call re-order labels on take_2d_labels_or_positional result if necessary.

        Parameters
        ----------
        intermediate : PandasDataFrame
        row_positions : list-like of ints, optional
            The row positions to extract.
        sorted_row_positions : list-like of ints, optional
            Sorted version of row_positions.
        col_positions : list-like of ints, optional
            The column positions to extract.
        sorted_col_positions : list-like of ints, optional
            Sorted version of col_positions.

        Returns
        -------
        PandasDataframe
        """
        # Check if monotonically increasing, return if it is. Fast track code path for
        # common case to keep it fast.
        if (
            row_positions is None
            # Fast range processing of non-positive-step ranges is not yet supported
            or (is_range_like(row_positions) and row_positions.step > 0)
            or len(row_positions) == 1
            or n.all(row_positions[1:] >= row_positions[:-1])
        ) and (
            col_positions is None
            # Fast range processing of non-positive-step ranges is not yet supported
            or (is_range_like(col_positions) and col_positions.step > 0)
            or len(col_positions) == 1
            or n.all(col_positions[1:] >= col_positions[:-1])
        ):
            return intermediate

        # The new labels are often smaller than the old labels, so we can't reuse the
        # original order values because those were mapped to the original data. We have
        # to reorder here based on the expected order from within the data.
        # We create a dictionary mapping the position of the numeric index with respect
        # to all others, then recreate that order by mapping the new order values from
        # the old. This information is sent to `_reorder_labels`.
        if row_positions is not None:
            row_order_mapping = dict(
                zip(sorted_row_positions, range(len(row_positions)))
            )
            new_row_order = [row_order_mapping[idx] for idx in row_positions]
        else:
            new_row_order = None
        if col_positions is not None:
            col_order_mapping = dict(
                zip(sorted_col_positions, range(len(col_positions)))
            )
            new_col_order = [col_order_mapping[idx] for idx in col_positions]
        else:
            new_col_order = None
        return intermediate._reorder_labels(
            row_positions=new_row_order, col_positions=new_col_order
        )
//  [python]
def _maybe_reorder_labels(
        self,
        intermediate: "PandasDataframe",
        row_positions,
        sorted_row_positions,
        col_positions,
        sorted_col_positions,
    ) -> "PandasDataframe":
        """
        Call re-order labels on take_2d_labels_or_positional result if necessary.

        Parameters
        ----------
        intermediate : PandasDataFrame
        row_positions : list-like of ints, optional
            The row positions to extract.
        sorted_row_positions : list-like of ints, optional
            Sorted version of row_positions.
        col_positions : list-like of ints, optional
            The column positions to extract.
        sorted_col_positions : list-like of ints, optional
            Sorted version of col_positions.

        Returns
        -------
        PandasDataframe
        """
        # Check if monotonically increasing, return if it is. Fast track code path for
        # common case to keep it fast.
        if (
            row_positions is None
            # Fast range processing of non-positive-step ranges is not yet supported
            or (is_range_like(row_positions) and row_positions.step > 0)
            or len(row_positions) == 1
            or np.all(row_positions[1:] >= row_positions[:-1])
        ) and (
            col_positions is None
            # Fast range processing of non-positive-step ranges is not yet supported
            or (is_range_like(col_positions) and col_positions.step > 0)
            or len(col_positions) == 1
            or np.all(col_positions[1:] >= col_positions[:-1])
        ):
            return intermediate

        # The new labels are often smaller than the old labels, so we can't reuse the
        # original order values because those were mapped to the original data. We have
        # to reorder here based on the expected order from within the data.
        # We create a dictionary mapping the position of the numeric index with respect
        # to all others, then recreate that order by mapping the new order values from
        # the old. This information is sent to `_reorder_labels`.
        if row_positions is not None:
            row_order_mapping = dict(
                zip(sorted_row_positions, range(len(row_positions)))
            )
            new_row_order = [row_order_mapping[idx] for idx in row_positions]
        else:
            new_row_order = None
        if col_positions is not None:
            col_order_mapping = dict(
                zip(sorted_col_positions, range(len(col_positions)))
            )
            new_col_order = [col_order_mapping[idx] for idx in col_positions]
        else:
            new_col_order = None
        return intermediate._reorder_labels(
            row_positions=new_row_order, col_positions=new_col_order
        )
//  [python]
def local_neg_div_neg(node):
    """
    - (-a / b) -> a / b

    Also performs - (c / b) -> ((-c) / b) when c is a scalar constant.

    """
    if node.op == T.neg:
        if node.inputs[0].owner and node.inputs[0].owner.op == T.true_div:
            frac = node.inputs[0]
            num, denom = frac.owner.inputs
            if num.owner and num.owner.op == T.neg:
                if len(frac.clients) == 1:
                    # No other clients of the original division
                    new_num = num.owner.inputs[0]
                    return [T.true_div(new_num, denom)]
            elif n.all(num.broadcastable) and isinstance(num, Constant):
                if len(frac.clients) == 1:
                    new_num = -num.data
                    return [T.true_div(new_num, denom)]
//  [python]
def local_neg_div_neg(node):
    """
    - (-a / b) -> a / b

    Also performs - (c / b) -> ((-c) / b) when c is a scalar constant.

    """
    if node.op == T.neg:
        if node.inputs[0].owner and node.inputs[0].owner.op == T.true_div:
            frac = node.inputs[0]
            num, denom = frac.owner.inputs
            if num.owner and num.owner.op == T.neg:
                if len(frac.clients) == 1:
                    # No other clients of the original division
                    new_num = num.owner.inputs[0]
                    return [T.true_div(new_num, denom)]
            elif numpy.all(num.broadcastable) and isinstance(num, Constant):
                if len(frac.clients) == 1:
                    new_num = -num.data
                    return [T.true_div(new_num, denom)]
//  [python]
def eig(cls, harness: test_harnesses.Harness):
    compute_left_eigenvectors = harness.params["compute_left_eigenvectors"]
    compute_right_eigenvectors = harness.params["compute_right_eigenvectors"]
    dtype = harness.dtype

    def custom_assert(tst, result_jax, result_tf, *, args, tol, err_msg):
      operand, = args
      inner_dimension = operand.shape[-1]

      # Test ported from tests.linlag_test.testEig
      # Norm, adjusted for dimension and type.
      def norm(x):
        norm = np.linalg.norm(x, axis=(-2, -1))
        return norm / ((inner_dimension + 1) * jnp.finfo(dtype).eps)

      def check_right_eigenvectors(a, w, vr):
        tst.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)

      def check_eigenvalue_is_in_array(eigenvalue, eigenvalues_array):
        tol = None
        # TODO(bchetioui): numerical discrepancies
        if dtype in [np.float32, np.complex64]:
          tol = 1e-4
        elif dtype in [np.float64, np.complex128]:
          tol = 1e-13
        closest_diff = min(abs(eigenvalues_array - eigenvalue))
        tst.assertAllClose(
            closest_diff,
            np.array(0., closest_diff.dtype),
            atol=tol,
            err_msg=err_msg)

      all_w_jax, all_w_tf = result_jax[0], result_tf[0]
      for idx in itertools.product(*map(range, operand.shape[:-2])):
        w_jax, w_tf = all_w_jax[idx], all_w_tf[idx]
        for i in range(inner_dimension):
          check_eigenvalue_is_in_array(w_jax[i], w_tf)
          check_eigenvalue_is_in_array(w_tf[i], w_jax)

      if compute_left_eigenvectors:
        check_left_eigenvectors(operand, all_w_tf, result_tf[1])
      if compute_right_eigenvectors:
        check_right_eigenvectors(operand, all_w_tf,
                                 result_tf[1 + compute_left_eigenvectors])

    return [
        # Eig does not work in JAX on gpu or tpu
        Jax2TfLimitation(
            "function not compilable", modes="compiled", devices="cpu"),
        Jax2TfLimitation(
            "TF Conversion of eig is not implemented when both compute_left_eigenvectors and compute_right_eigenvectors are set to True",
            enabled=(compute_left_eigenvectors and compute_right_eigenvectors)),
        custom_numeric(
            custom_assert=custom_assert,
            description=("May return the eigenvalues and eigenvectors in a "
                         "potentially different order. The eigenvectors may "
                         "also be different, but equally valid."))
    ]
//  [python]
def eig(cls, harness: test_harnesses.Harness):
    compute_left_eigenvectors = harness.params["compute_left_eigenvectors"]
    compute_right_eigenvectors = harness.params["compute_right_eigenvectors"]
    dtype = harness.dtype

    def custom_assert(tst, result_jax, result_tf, *, args, tol, err_msg):
      operand, = args
      inner_dimension = operand.shape[-1]

      # Test ported from tests.linlag_test.testEig
      # Norm, adjusted for dimension and type.
      def norm(x):
        norm = np.linalg.norm(x, axis=(-2, -1))
        return norm / ((inner_dimension + 1) * jnp.finfo(dtype).eps)

      def check_right_eigenvectors(a, w, vr):
        tst.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)

      def check_eigenvalue_is_in_array(eigenvalue, eigenvalues_array):
        tol = None
        # TODO(bchetioui): numerical discrepancies
        if dtype in [np.float32, np.complex64]:
          tol = 1e-4
        elif dtype in [np.float64, np.complex128]:
          tol = 1e-13
        closest_diff = min(abs(eigenvalues_array - eigenvalue))
        tst.assertAllClose(
            closest_diff,
            np.array(0., closest_diff.dtype),
            atol=tol,
            err_msg=err_msg)

      all_w_jax, all_w_tf = result_jax[0], result_tf[0]
      for idx in itertools.product(*map(range, operand.shape[:-2])):
        w_jax, w_tf = all_w_jax[idx], all_w_tf[idx]
        for i in range(inner_dimension):
          check_eigenvalue_is_in_array(w_jax[i], w_tf)
          check_eigenvalue_is_in_array(w_tf[i], w_jax)

      if compute_left_eigenvectors:
        check_left_eigenvectors(operand, all_w_tf, result_tf[1])
      if compute_right_eigenvectors:
        check_right_eigenvectors(operand, all_w_tf,
                                 result_tf[1 + compute_left_eigenvectors])

    return [
        # Eig does not work in JAX on gpu or tpu
        Jax2TfLimitation(
            "function not compilable", modes="compiled", devices="cpu"),
        Jax2TfLimitation(
            "TF Conversion of eig is not implemented when both compute_left_eigenvectors and compute_right_eigenvectors are set to True",
            enabled=(compute_left_eigenvectors and compute_right_eigenvectors)),
        custom_numeric(
            custom_assert=custom_assert,
            description=("May return the eigenvalues and eigenvectors in a "
                         "potentially different order. The eigenvectors may "
                         "also be different, but equally valid."))
    ]
//  [python]
def local_div_to_inv(node):
    if node.op == T.true_div and n.all(
            local_mul_canonizer.get_constant(node.inputs[0]) == 1.0):
        out = node.outputs[0]
        new_out = T.inv(local_mul_canonizer.merge_num_denum(node.inputs[1:],
                                                            []))
        # The ones could have forced upcasting
        if new_out.dtype != out.dtype:
            new_out = T.cast(new_out, dtype=out.dtype)
        # The ones could have forced a specific length
        if new_out.type != out.type:
            new_out = broadcast_like(new_out, out, node.fgraph)
        return [new_out]
    else:
        return False
//  [python]
def local_div_to_inv(node):
    if node.op == T.true_div and numpy.all(
            local_mul_canonizer.get_constant(node.inputs[0]) == 1.0):
        out = node.outputs[0]
        new_out = T.inv(local_mul_canonizer.merge_num_denum(node.inputs[1:],
                                                            []))
        # The ones could have forced upcasting
        if new_out.dtype != out.dtype:
            new_out = T.cast(new_out, dtype=out.dtype)
        # The ones could have forced a specific length
        if new_out.type != out.type:
            new_out = broadcast_like(new_out, out, node.fgraph)
        return [new_out]
    else:
        return False
//  [python]
def local_intdiv_by_one(node):
    """x // 1 -> x
    """
    if node.op in [T.int_div]:
        if isinstance(node.inputs[1], T.TensorConstant) and \
           n.all(node.inputs[1].value == 1):
            return [node.inputs[0].astype(node.outputs[0].dtype)]
//  [python]
def local_intdiv_by_one(node):
    """x // 1 -> x
    """
    if node.op in [T.int_div]:
        if isinstance(node.inputs[1], T.TensorConstant) and \
           numpy.all(node.inputs[1].value == 1):
            return [node.inputs[0].astype(node.outputs[0].dtype)]
//  [python]
def add_calculate(num, denum, aslist=False, out_type=None):
    if out_type is None:
        zero = 0.0
    else:
        zero = theano._asarray(0, dtype=out_type.dtype)
    v = reduce(numpy.add, num, zero) - reduce(numpy.add, denum, zero)
    if aslist:
        if numpy.all(v == 0):
            return []
        else:
            return [v]
    return v
//  [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(numpy.add, num, zero) - reduce(numpy.add, denum, zero)
    if aslist:
        if numpy.all(v == 0):
            return []
        else:
            return [v]
    return v
//  [python]
def transform(self, node):
        op = node.op
        if op not in [self.main, self.inverse, self.reciprocal]:
            return False

        assert len(node.outputs) == 1
        out = node.outputs[0]

        # check if any of the clients of this node would be part of
        # this canonized graph...  if so, we do nothing and wait for
        # them to be transformed.
        for c, c_idx in out.clients:
            if c == 'output':
                continue
            while (isinstance(getattr(c, 'op', None), DimShuffle) and
                   len(c.outputs[0].clients) <= 1):
                c = c.outputs[0].clients[0][0]
            if getattr(c, 'op', '') in [self.main, self.inverse,
                                        self.reciprocal]:
                return False

        # Here we make the canonical version of the graph around this node
        # See the documentation of get_num_denum and simplify
        orig_num, orig_denum = self.get_num_denum(node.outputs[0])
        num, denum = self.simplify(list(orig_num), list(orig_denum), out.type)

        def same(x, y):
            return len(x) == len(y) and all(n.all(xe == ye) for xe, ye in
                                            zip(x, y))

        if same(orig_num, num) and same(orig_denum, denum):
            # We return False if there are no changes
            return False

        new = self.merge_num_denum(num, denum)
        if new.type.dtype != out.type.dtype:
            # new = T.fill(out, new)
            elem_op = T.Elemwise(scalar.Identity(scalar.specific_out(
                getattr(scalar, out.type.dtype))))
            new = elem_op(new)

        assert (new.type == out.type) == (not (new.type != out.type))

        if not (new.type == out.type):
            new = _fill_chain(new, node.inputs)[0]

        if new.type == out.type:
            # This happen with test
            # theano/tensor/tests/test_opt.py:T_local_switch_sink
            new.tag.values_eq_approx = values_eq_approx_remove_inf_nan
            return [new]
        else:
            _logger.warning(' '.join(('CANONIZE FAILED: new, out = ',
                                      new, ',', out, 'types',
                                      new.type, ',', out.type)))
            return False
//  [python]
def transform(self, node):
        op = node.op
        if op not in [self.main, self.inverse, self.reciprocal]:
            return False

        assert len(node.outputs) == 1
        out = node.outputs[0]

        # check if any of the clients of this node would be part of
        # this canonized graph...  if so, we do nothing and wait for
        # them to be transformed.
        for c, c_idx in out.clients:
            if c == 'output':
                continue
            while (isinstance(getattr(c, 'op', None), DimShuffle) and
                   len(c.outputs[0].clients) <= 1):
                c = c.outputs[0].clients[0][0]
            if getattr(c, 'op', '') in [self.main, self.inverse,
                                        self.reciprocal]:
                return False

        # Here we make the canonical version of the graph around this node
        # See the documentation of get_num_denum and simplify
        orig_num, orig_denum = self.get_num_denum(node.outputs[0])
        num, denum = self.simplify(list(orig_num), list(orig_denum), out.type)

        def same(x, y):
            return len(x) == len(y) and all(numpy.all(xe == ye) for xe, ye in
                                            zip(x, y))

        if same(orig_num, num) and same(orig_denum, denum):
            # We return False if there are no changes
            return False

        new = self.merge_num_denum(num, denum)
        if new.type.dtype != out.type.dtype:
            # new = T.fill(out, new)
            elem_op = T.Elemwise(scalar.Identity(scalar.specific_out(
                getattr(scalar, out.type.dtype))))
            new = elem_op(new)

        assert (new.type == out.type) == (not (new.type != out.type))

        if not (new.type == out.type):
            new = _fill_chain(new, node.inputs)[0]

        if new.type == out.type:
            # This happen with test
            # theano/tensor/tests/test_opt.py:T_local_switch_sink
            new.tag.values_eq_approx = values_eq_approx_remove_inf_nan
            return [new]
        else:
            _logger.warning(' '.join(('CANONIZE FAILED: new, out = ',
                                      new, ',', out, 'types',
                                      new.type, ',', out.type)))
            return False
//  [python]
def local_sum_prod_mul_by_scalar(node):
    """
    sum(scalar * smth) -> scalar * sum(smth)
    sum(-smth) -> -sum(smth)

    or

    prod(scalar * smth) -> scalar ** size(smth) * prod(smth)
    prod(-smth) -> -1 ** size(smth) * prod(smth)

    """
    # TODO: if the the thing inside the Sum is a division,
    # we should get at the numerator....
    if isinstance(node.op, (T.Sum, T.elemwise.Prod)):
        node_inps, = node.inputs
        if node_inps.owner and node_inps.owner.op == T.mul:
            terms = node_inps.owner.inputs
            scalars = [t.dimshuffle() for t in terms if
                       n.all(t.type.broadcastable)]

            if len(scalars) == 0:
                # Nothing to optimize here
                return

            non_scalars = [t for t in terms if not n.all(t.broadcastable)]

            # Perform the op only on the non-scalar inputs, if applicable
            if len(non_scalars) == 0:
                new_op_input_nb_elements = 1
                new_op_output = 1
            elif len(non_scalars) == 1:
                new_op_input_nb_elements = non_scalars[0].size
                new_op_output = node.op(non_scalars[0])
            else:
                new_op_input = T.mul(*non_scalars)
                new_op_input_nb_elements = new_op_input.size
                new_op_output = node.op(new_op_input)

            # If node.op is a T.elemwise.Prod, then the scalars need to be
            # raised to the power of the number of elements in the input
            # to the Prod
            if (isinstance(node.op, T.elemwise.Prod) and
                    new_op_input_nb_elements != 1):

                scalars = [s ** new_op_input_nb_elements for s in scalars]

            # Scale the output of the op by the scalars and return as
            # replacement for the original output
            mul_inputs = scalars
            if new_op_input_nb_elements != 1:
                mul_inputs.append(new_op_output)

            if len(mul_inputs) == 1:
                return mul_inputs
            else:
                return [T.mul(*mul_inputs)]

        if isinstance(node.op, T.Sum) and node_inps.owner and node_inps.owner.op == T.neg:
            return [T.neg(node.op(node_inps.owner.inputs[0]))]
//  [python]
def local_sum_prod_mul_by_scalar(node):
    """
    sum(scalar * smth) -> scalar * sum(smth)
    sum(-smth) -> -sum(smth)

    or

    prod(scalar * smth) -> scalar ** size(smth) * prod(smth)
    prod(-smth) -> -1 ** size(smth) * prod(smth)

    """
    # TODO: if the the thing inside the Sum is a division,
    # we should get at the numerator....
    if isinstance(node.op, (T.Sum, T.elemwise.Prod)):
        node_inps, = node.inputs
        if node_inps.owner and node_inps.owner.op == T.mul:
            terms = node_inps.owner.inputs
            scalars = [t.dimshuffle() for t in terms if
                       numpy.all(t.type.broadcastable)]

            if len(scalars) == 0:
                # Nothing to optimize here
                return

            non_scalars = [t for t in terms if not numpy.all(t.broadcastable)]

            # Perform the op only on the non-scalar inputs, if applicable
            if len(non_scalars) == 0:
                new_op_input_nb_elements = 1
                new_op_output = 1
            elif len(non_scalars) == 1:
                new_op_input_nb_elements = non_scalars[0].size
                new_op_output = node.op(non_scalars[0])
            else:
                new_op_input = T.mul(*non_scalars)
                new_op_input_nb_elements = new_op_input.size
                new_op_output = node.op(new_op_input)

            # If node.op is a T.elemwise.Prod, then the scalars need to be
            # raised to the power of the number of elements in the input
            # to the Prod
            if (isinstance(node.op, T.elemwise.Prod) and
                    new_op_input_nb_elements != 1):

                scalars = [s ** new_op_input_nb_elements for s in scalars]

            # Scale the output of the op by the scalars and return as
            # replacement for the original output
            mul_inputs = scalars
            if new_op_input_nb_elements != 1:
                mul_inputs.append(new_op_output)

            if len(mul_inputs) == 1:
                return mul_inputs
            else:
                return [T.mul(*mul_inputs)]

        if isinstance(node.op, T.Sum) and node_inps.owner and node_inps.owner.op == T.neg:
            return [T.neg(node.op(node_inps.owner.inputs[0]))]
//  [python]
def local_pow_specialize(node):
    # here, we are past the point of canonicalization, so we don't want
    # to put in un-necessary fills.
    if node.op == T.pow:
        # the idea here is that we have pow(x, y)
        odtype = node.outputs[0].dtype
        xsym = node.inputs[0]
        ysym = node.inputs[1]
        y = local_mul_canonizer.get_constant(ysym)
        if (y is not None) \
                and encompasses_broadcastable(xsym.type.broadcastable,
                                              ysym.type.broadcastable):
            rval = None

            if n.all(y == 2):
                rval = [T.sqr(xsym)]
            if n.all(y == 1):
                rval = [xsym]
            if n.all(y == 0):
                rval = [T.fill(xsym, numpy.asarray(1, dtype=odtype))]
            if n.all(y == 0.5):
                rval = [T.sqrt(xsym)]
            if n.all(y == -0.5):
                rval = [T.inv(T.sqrt(xsym))]
            if n.all(y == -1):
                rval = [T.inv(xsym)]
            if n.all(y == -2):
                rval = [T.inv(T.sqr(xsym))]
            if rval:
                rval[0] = T.cast(rval[0], odtype)
                assert rval[0].type == node.outputs[0].type, (
                    rval, node.outputs)
                return rval
    else:
        return False
//  [python]
def local_pow_specialize(node):
    # here, we are past the point of canonicalization, so we don't want
    # to put in un-necessary fills.
    if node.op == T.pow:
        # the idea here is that we have pow(x, y)
        odtype = node.outputs[0].dtype
        xsym = node.inputs[0]
        ysym = node.inputs[1]
        y = local_mul_canonizer.get_constant(ysym)
        if (y is not None) \
                and encompasses_broadcastable(xsym.type.broadcastable,
                                              ysym.type.broadcastable):
            rval = None

            if numpy.all(y == 2):
                rval = [T.sqr(xsym)]
            if numpy.all(y == 1):
                rval = [xsym]
            if numpy.all(y == 0):
                rval = [T.fill(xsym, numpy.asarray(1, dtype=odtype))]
            if numpy.all(y == 0.5):
                rval = [T.sqrt(xsym)]
            if numpy.all(y == -0.5):
                rval = [T.inv(T.sqrt(xsym))]
            if numpy.all(y == -1):
                rval = [T.inv(xsym)]
            if numpy.all(y == -2):
                rval = [T.inv(T.sqr(xsym))]
            if rval:
                rval[0] = T.cast(rval[0], odtype)
                assert rval[0].type == node.outputs[0].type, (
                    rval, node.outputs)
                return rval
    else:
        return False