jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_stable_attention_weights():
    state_dim = 2
    attended_dim = 3
    match_dim = 4
    attended_length = 5
    batch_size = 6

    attention = SequenceContentAttention(
        state_names=["states"], state_dims=[state_dim],
        attended_dim=attended_dim, match_dim=match_dim,
        weights_init=IsotropicGaussian(0.5),
        biases_init=Constant(0))
    attention.initialize()

    # Random high energies with mu=800, sigma=50
    energies_val = (
        50. * numpy.random.randn(attended_length, batch_size) + 800
        ).astype(theano.config.floatX)
    energies = tensor.as_tensor_variable(energies_val)
    mask = tensor.as_tensor_variable(
        numpy.ones((attended_length, batch_size)))
    weights = attention.compute_weights(energies, mask).eval()
    assert n.all(n.isfinite(weights))
//  [python]
def test_stable_attention_weights():
    state_dim = 2
    attended_dim = 3
    match_dim = 4
    attended_length = 5
    batch_size = 6

    attention = SequenceContentAttention(
        state_names=["states"], state_dims=[state_dim],
        attended_dim=attended_dim, match_dim=match_dim,
        weights_init=IsotropicGaussian(0.5),
        biases_init=Constant(0))
    attention.initialize()

    # Random high energies with mu=800, sigma=50
    energies_val = (
        50. * numpy.random.randn(attended_length, batch_size) + 800
        ).astype(theano.config.floatX)
    energies = tensor.as_tensor_variable(energies_val)
    mask = tensor.as_tensor_variable(
        numpy.ones((attended_length, batch_size)))
    weights = attention.compute_weights(energies, mask).eval()
    assert numpy.all(numpy.isfinite(weights))
//  [python]
def add_brute_force(self, ps):
        relations = ps.relations
        if self.stop_on_collide:
            mask = np.logical_and(relations.dist < self.max_distance, relations.attract_mask)
        else:
            mask = relations.dist < self.max_distance
        index_non_inter = ps.relations.indexes[mask]
        id0 = index_non_inter[:, 0]
        id1 = index_non_inter[:, 1]
        dist2 = np.power(relations.dist[mask], self.decay)[:, np.newaxis]
        normal = relations.normal_v[mask]
        direction = normal / dist2 * relations.mass_product[mask, np.newaxis]


        att = self.f_magnitude
        len0, len1 = [att, att] if self.uniform_magnitude else [att[id1], att[id0]]

        np.add.at(ps.force_resultant, id0, -direction * len0)
        np.add.at(ps.force_resultant, id1, direction * len1)
//  [python]
def add_brute_force(self, ps):
        relations = ps.relations
        if self.stop_on_collide:
            mask = np.all((relations.dist < self.max_distance, relations.attract_mask), axis=0)
        else:
            mask = relations.dist < self.max_distance
        index_non_inter = ps.relations.indexes[mask]
        id0 = index_non_inter[:, 0]
        id1 = index_non_inter[:, 1]
        dist2 = np.power(relations.dist[mask], self.decay)[:, np.newaxis]
        normal = relations.normal_v[mask]
        direction = normal / dist2 * relations.mass_product[mask, np.newaxis]


        att = self.f_magnitude
        len0, len1 = [att, att] if self.uniform_magnitude else [att[id1], att[id0]]

        np.add.at(ps.force_resultant, id0, -direction * len0)
        np.add.at(ps.force_resultant, id1, direction * len1)
//  [python]
def _make_eigvals(alpha, beta, homogeneous_eigvals):
    if homogeneous_eigvals:
        if beta is None:
            return numpy.vstack((alpha, numpy.ones_like(alpha)))
        else:
            return numpy.vstack((alpha, beta))
    else:
        if beta is None:
            return alpha
        else:
            w = numpy.empty_like(alpha)
            alpha_zero = (alpha == 0)
            beta_zero = (beta == 0)
            beta_nonzero = ~beta_zero
            w[beta_nonzero] = alpha[beta_nonzero]/beta[beta_nonzero]
            # Use numpy.inf for complex values too since
            # 1/numpy.inf = 0, i.e. it correctly behaves as projective
            # infinity.
            w[~alpha_zero & beta_zero] = numpy.inf
            if n.all(alpha.imag == 0):
                w[alpha_zero & beta_zero] = numpy.nan
            else:
                w[alpha_zero & beta_zero] = complex(numpy.nan, numpy.nan)
            return w
//  [python]
def _make_eigvals(alpha, beta, homogeneous_eigvals):
    if homogeneous_eigvals:
        if beta is None:
            return numpy.vstack((alpha, numpy.ones_like(alpha)))
        else:
            return numpy.vstack((alpha, beta))
    else:
        if beta is None:
            return alpha
        else:
            w = numpy.empty_like(alpha)
            alpha_zero = (alpha == 0)
            beta_zero = (beta == 0)
            beta_nonzero = ~beta_zero
            w[beta_nonzero] = alpha[beta_nonzero]/beta[beta_nonzero]
            # Use numpy.inf for complex values too since
            # 1/numpy.inf = 0, i.e. it correctly behaves as projective
            # infinity.
            w[~alpha_zero & beta_zero] = numpy.inf
            if numpy.all(alpha.imag == 0):
                w[alpha_zero & beta_zero] = numpy.nan
            else:
                w[alpha_zero & beta_zero] = complex(numpy.nan, numpy.nan)
            return w
//  [python]
def zpk2tf(z, p, k):
    """
    Return polynomial transfer function representation from zeros and poles

    Parameters
    ----------
    z : array_like
        Zeros of the transfer function.
    p : array_like
        Poles of the transfer function.
    k : float
        System gain.

    Returns
    -------
    b : ndarray
        Numerator polynomial coefficients.
    a : ndarray
        Denominator polynomial coefficients.

    """
    z = atleast_1d(z)
    k = atleast_1d(k)
    if len(z.shape) > 1:
        temp = poly(z[0])
        b = zeros((z.shape[0], z.shape[1] + 1), temp.dtype.char)
        if len(k) == 1:
            k = [k[0]] * z.shape[0]
        for i in range(z.shape[0]):
            b[i] = k[i] * poly(z[i])
    else:
        b = k * poly(z)
    a = atleast_1d(poly(p))

    # Use real output if possible.  Copied from numpy.poly, since
    # we can't depend on a specific version of numpy.
    if issubclass(b.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(z, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.alltrue(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                b = b.real.copy()

    if issubclass(a.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(p, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.alltrue(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                a = a.real.copy()

    return b, a
//  [python]
def zpk2tf(z, p, k):
    """
    Return polynomial transfer function representation from zeros and poles

    Parameters
    ----------
    z : array_like
        Zeros of the transfer function.
    p : array_like
        Poles of the transfer function.
    k : float
        System gain.

    Returns
    -------
    b : ndarray
        Numerator polynomial coefficients.
    a : ndarray
        Denominator polynomial coefficients.

    """
    z = atleast_1d(z)
    k = atleast_1d(k)
    if len(z.shape) > 1:
        temp = poly(z[0])
        b = zeros((z.shape[0], z.shape[1] + 1), temp.dtype.char)
        if len(k) == 1:
            k = [k[0]] * z.shape[0]
        for i in range(z.shape[0]):
            b[i] = k[i] * poly(z[i])
    else:
        b = k * poly(z)
    a = atleast_1d(poly(p))

    # Use real output if possible.  Copied from numpy.poly, since
    # we can't depend on a specific version of numpy.
    if issubclass(b.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(z, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.all(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                b = b.real.copy()

    if issubclass(a.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(p, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.all(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                a = a.real.copy()

    return b, a
//  [python]
def normalize(b, a):
    """Normalize numerator/denominator of a continuous-time transfer function.

    If values of `b` are too close to 0, they are removed. In that case, a
    BadCoefficients warning is emitted.

    Parameters
    ----------
    b: array_like
        Numerator of the transfer function. Can be a 2d array to normalize
        multiple transfer functions.
    a: array_like
        Denominator of the transfer function. At most 1d.

    Returns
    -------
    num: array
        The numerator of the normalized transfer function. At least a 1d
        array. A 2d-array if the input `num` is a 2d array.
    den: 1d-array
        The denominator of the normalized transfer function.

    Notes
    -----
    Coefficients for both the numerator and denominator should be specified in
    descending exponent order (e.g., ``s^2 + 3s + 5`` would be represented as
    ``[1, 3, 5]``).
    """
    num, den = b, a

    den = np.atleast_1d(den)
    num = np.atleast_2d(_align_nums(num))

    if den.ndim != 1:
        raise ValueError("Denominator polynomial must be rank-1 array.")
    if num.ndim > 2:
        raise ValueError("Numerator polynomial must be rank-1 or"
                         " rank-2 array.")
    if np.all(den == 0):
        raise ValueError("Denominator must have at least on nonzero element.")

    # Trim leading zeros in denominator, leave at least one.
    den = np.trim_zeros(den, 'f')

    # Normalize transfer function
    num, den = num / den[0], den / den[0]

    # Count numerator columns that are all zero
    leading_zeros = 0
    for col in num.T:
        if np.allclose(col, 0, atol=1e-14):
            leading_zeros += 1
        else:
            break

    # Trim leading zeros of numerator
    if leading_zeros > 0:
        warnings.warn("Badly conditioned filter coefficients (numerator): the "
                      "results may be meaningless", BadCoefficients)
        # Make sure at least one column remains
        if leading_zeros == num.shape[1]:
            leading_zeros -= 1
        num = num[:, leading_zeros:]

    # Squeeze first dimension if singular
    if num.shape[0] == 1:
        num = num[0, :]

    return num, den
//  [python]
def values_eq(self, a, b):
        sa = a.get_state()
        sb = b.get_state()
        # Should always be the string 'MT19937'
        if sa[0] != sb[0]:
            return False
        # 1-D array of 624 unsigned integer keys
        if not n.all(sa[1] == sb[1]):
            return False
        # integer "pos" representing the position in the array
        if sa[2] != sb[2]:
            return False
        # integer "has_gauss"
        if sa[3] != sb[3]:
            return False
        # float "cached_gaussian".
        # /!\ It is not initialized if has_gauss == 0
        if sa[3] != 0:
            if sa[4] != sb[4]:
                return False
        return True
//  [python]
def values_eq(self, a, b):
        sa = a.get_state()
        sb = b.get_state()
        # Should always be the string 'MT19937'
        if sa[0] != sb[0]:
            return False
        # 1-D array of 624 unsigned integer keys
        if not numpy.all(sa[1] == sb[1]):
            return False
        # integer "pos" representing the position in the array
        if sa[2] != sb[2]:
            return False
        # integer "has_gauss"
        if sa[3] != sb[3]:
            return False
        # float "cached_gaussian".
        # /!\ It is not initialized if has_gauss == 0
        if sa[3] != 0:
            if sa[4] != sb[4]:
                return False
        return True
//  [python]
def test_03_02_empty_mean(self):
        # Regression test - if there are no children, the per-parent means
        #                   should still be populated
        i, j = numpy.mgrid[0:20, 0:20]
        parent_labels = (i / 10 + 1).astype(int)
        child_labels = numpy.zeros(parent_labels.shape, int)
        workspace, module = self.make_workspace(
            parent_labels, child_labels, fake_measurement=True
        )
        module.wants_per_parent_means.value = True
        m = workspace.measurements
        self.assertTrue(isinstance(m,cellprofiler_core.measurement.Measurements))
        m.add_measurement(CHILD_OBJECTS, MEASUREMENT, numpy.zeros(0))
        m.add_measurement(CHILD_OBJECTS, IGNORED_MEASUREMENT, numpy.zeros(0, int))
        module.run(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, MEASUREMENT)
        self.assertTrue(name in m.get_feature_names(PARENT_OBJECTS))
        data = m.get_current_measurement(PARENT_OBJECTS, name)
        self.assertTrue(numpy.isnan(data).all())
        self.features_and_columns_match(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, IGNORED_MEASUREMENT)
        self.assertFalse(name in m.get_feature_names(PARENT_OBJECTS))
//  [python]
def test_03_02_empty_mean(self):
        # Regression test - if there are no children, the per-parent means
        #                   should still be populated
        i, j = numpy.mgrid[0:20, 0:20]
        parent_labels = (i / 10 + 1).astype(int)
        child_labels = numpy.zeros(parent_labels.shape, int)
        workspace, module = self.make_workspace(
            parent_labels, child_labels, fake_measurement=True
        )
        module.wants_per_parent_means.value = True
        m = workspace.measurements
        self.assertTrue(isinstance(m,cellprofiler_core.measurement.Measurements))
        m.add_measurement(CHILD_OBJECTS, MEASUREMENT, numpy.zeros(0))
        m.add_measurement(CHILD_OBJECTS, IGNORED_MEASUREMENT, numpy.zeros(0, int))
        module.run(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, MEASUREMENT)
        self.assertTrue(name in m.get_feature_names(PARENT_OBJECTS))
        data = m.get_current_measurement(PARENT_OBJECTS, name)
        self.assertTrue(numpy.all(numpy.isnan(data)))
        self.features_and_columns_match(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, IGNORED_MEASUREMENT)
        self.assertFalse(name in m.get_feature_names(PARENT_OBJECTS))
//  [python]
def test_03_01_mean(self):
        """Compute the mean for two parents and four children"""
        i, j = numpy.mgrid[0:20, 0:20]
        parent_labels = (i / 10 + 1).astype(int)
        child_labels = (i / 10).astype(int) + (j / 10).astype(int) * 2 + 1
        workspace, module = self.make_workspace(
            parent_labels, child_labels, fake_measurement=True
        )
        module.wants_per_parent_means.value = True
        m = workspace.measurements
        self.assertTrue(isinstance(m,cellprofiler_core.measurement.Measurements))
        m.add_measurement(CHILD_OBJECTS, MEASUREMENT, numpy.array([1.0, 2.0, 3.0, 4.0]))
        m.add_measurement(CHILD_OBJECTS, IGNORED_MEASUREMENT, numpy.array([1, 2, 3, 4]))
        expected = numpy.array([2.0, 3.0])
        module.run(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, MEASUREMENT)
        self.assertTrue(name in m.get_feature_names(PARENT_OBJECTS))
        data = m.get_current_measurement(PARENT_OBJECTS, name)
        self.assertTrue(numpy.alltrue(data == expected))
        self.features_and_columns_match(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, IGNORED_MEASUREMENT)
        self.assertFalse(name in m.get_feature_names(PARENT_OBJECTS))
//  [python]
def test_03_01_mean(self):
        """Compute the mean for two parents and four children"""
        i, j = numpy.mgrid[0:20, 0:20]
        parent_labels = (i / 10 + 1).astype(int)
        child_labels = (i / 10).astype(int) + (j / 10).astype(int) * 2 + 1
        workspace, module = self.make_workspace(
            parent_labels, child_labels, fake_measurement=True
        )
        module.wants_per_parent_means.value = True
        m = workspace.measurements
        self.assertTrue(isinstance(m,cellprofiler_core.measurement.Measurements))
        m.add_measurement(CHILD_OBJECTS, MEASUREMENT, numpy.array([1.0, 2.0, 3.0, 4.0]))
        m.add_measurement(CHILD_OBJECTS, IGNORED_MEASUREMENT, numpy.array([1, 2, 3, 4]))
        expected = numpy.array([2.0, 3.0])
        module.run(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, MEASUREMENT)
        self.assertTrue(name in m.get_feature_names(PARENT_OBJECTS))
        data = m.get_current_measurement(PARENT_OBJECTS, name)
        self.assertTrue(numpy.all(data == expected))
        self.features_and_columns_match(workspace)
        name = "Mean_%s_%s" % (CHILD_OBJECTS, IGNORED_MEASUREMENT)
        self.assertFalse(name in m.get_feature_names(PARENT_OBJECTS))
//  [python]
def test_04_00_distance_empty(self):
        """Make sure we can handle labels matrices that are all zero"""
        empty_labels = numpy.zeros((10, 20), int)
        some_labels = numpy.zeros((10, 20), int)
        some_labels[2:7, 3:8] = 1
        some_labels[3:8, 12:17] = 2
        for parent_labels, child_labels, n in (
            (empty_labels, empty_labels, 0),
            (some_labels, empty_labels, 0),
            (empty_labels, some_labels, 2),
        ):
            workspace, module = self.make_workspace(parent_labels, child_labels)
            self.assertTrue(
                isinstance(module, cellprofiler.modules.relateobjects.Relate)
            )
            module.find_parent_child_distances.value = (
                cellprofiler.modules.relateobjects.D_BOTH
            )
            module.run(workspace)
            self.features_and_columns_match(workspace)
            meas = workspace.measurements
            for feature in (
                cellprofiler.modules.relateobjects.FF_CENTROID,
                cellprofiler.modules.relateobjects.FF_MINIMUM,
            ):
                m = feature % PARENT_OBJECTS
                v = meas.get_current_measurement(CHILD_OBJECTS, m)
                self.assertEqual(len(v), n)
                if n > 0:
                    self.assertTrue(n.all(numpy.isnan(v)))
//  [python]
def test_04_00_distance_empty(self):
        """Make sure we can handle labels matrices that are all zero"""
        empty_labels = numpy.zeros((10, 20), int)
        some_labels = numpy.zeros((10, 20), int)
        some_labels[2:7, 3:8] = 1
        some_labels[3:8, 12:17] = 2
        for parent_labels, child_labels, n in (
            (empty_labels, empty_labels, 0),
            (some_labels, empty_labels, 0),
            (empty_labels, some_labels, 2),
        ):
            workspace, module = self.make_workspace(parent_labels, child_labels)
            self.assertTrue(
                isinstance(module, cellprofiler.modules.relateobjects.Relate)
            )
            module.find_parent_child_distances.value = (
                cellprofiler.modules.relateobjects.D_BOTH
            )
            module.run(workspace)
            self.features_and_columns_match(workspace)
            meas = workspace.measurements
            for feature in (
                cellprofiler.modules.relateobjects.FF_CENTROID,
                cellprofiler.modules.relateobjects.FF_MINIMUM,
            ):
                m = feature % PARENT_OBJECTS
                v = meas.get_current_measurement(CHILD_OBJECTS, m)
                self.assertEqual(len(v), n)
                if n > 0:
                    self.assertTrue(numpy.all(numpy.isnan(v)))
//  [python]
def _geneig(a1, b1, left, right, overwrite_a, overwrite_b,
            homogeneous_eigvals):
    ggev, = get_lapack_funcs(('ggev',), (a1, b1))
    cvl, cvr = left, right
    res = ggev(a1, b1, lwork=-1)
    lwork = res[-2][0].real.astype(numpy.int)
    if ggev.typecode in 'cz':
        alpha, beta, vl, vr, work, info = ggev(a1, b1, cvl, cvr, lwork,
                                               overwrite_a, overwrite_b)
        w = _make_eigvals(alpha, beta, homogeneous_eigvals)
    else:
        alphar, alphai, beta, vl, vr, work, info = ggev(a1, b1, cvl, cvr,
                                                        lwork, overwrite_a,
                                                        overwrite_b)
        alpha = alphar + _I * alphai
        w = _make_eigvals(alpha, beta, homogeneous_eigvals)
    _check_info(info, 'generalized eig algorithm (ggev)')

    only_real = n.all(w.imag == 0.0)
    if not (ggev.typecode in 'cz' or only_real):
        t = w.dtype.char
        if left:
            vl = _make_complex_eigvecs(w, vl, t)
        if right:
            vr = _make_complex_eigvecs(w, vr, t)

    # the eigenvectors returned by the lapack function are NOT normalized
    for i in xrange(vr.shape[0]):
        if right:
            vr[:, i] /= norm(vr[:, i])
        if left:
            vl[:, i] /= norm(vl[:, i])

    if not (left or right):
        return w
    if left:
        if right:
            return w, vl, vr
        return w, vl
    return w, vr
//  [python]
def _geneig(a1, b1, left, right, overwrite_a, overwrite_b,
            homogeneous_eigvals):
    ggev, = get_lapack_funcs(('ggev',), (a1, b1))
    cvl, cvr = left, right
    res = ggev(a1, b1, lwork=-1)
    lwork = res[-2][0].real.astype(numpy.int)
    if ggev.typecode in 'cz':
        alpha, beta, vl, vr, work, info = ggev(a1, b1, cvl, cvr, lwork,
                                               overwrite_a, overwrite_b)
        w = _make_eigvals(alpha, beta, homogeneous_eigvals)
    else:
        alphar, alphai, beta, vl, vr, work, info = ggev(a1, b1, cvl, cvr,
                                                        lwork, overwrite_a,
                                                        overwrite_b)
        alpha = alphar + _I * alphai
        w = _make_eigvals(alpha, beta, homogeneous_eigvals)
    _check_info(info, 'generalized eig algorithm (ggev)')

    only_real = numpy.all(w.imag == 0.0)
    if not (ggev.typecode in 'cz' or only_real):
        t = w.dtype.char
        if left:
            vl = _make_complex_eigvecs(w, vl, t)
        if right:
            vr = _make_complex_eigvecs(w, vr, t)

    # the eigenvectors returned by the lapack function are NOT normalized
    for i in xrange(vr.shape[0]):
        if right:
            vr[:, i] /= norm(vr[:, i])
        if left:
            vl[:, i] /= norm(vl[:, i])

    if not (left or right):
        return w
    if left:
        if right:
            return w, vl, vr
        return w, vl
    return w, vr
//  [python]
def test_loc_assign(self) -> None:
        self.ds["x"] = ("x", np.array(list("abcdefghij")))
        da = self.ds["foo"]
        # assignment
        # typing issue: see https://github.com/python/mypy/issues/2410
        da.loc["a":"j"] = 0  # type: ignore[misc]
        assert n.all(da.values == 0)
        da.loc[{"x": slice("a", "j")}] = 2
        assert n.all(da.values == 2)

        da.loc[{"x": slice("a", "j")}] = 2
        assert n.all(da.values == 2)

        # Multi dimensional case
        da = DataArray(np.arange(12).reshape(3, 4), dims=["x", "y"])
        da.loc[0, 0] = 0
        assert da.values[0, 0] == 0
        assert da.values[0, 1] != 0

        da = DataArray(np.arange(12).reshape(3, 4), dims=["x", "y"])
        da.loc[0] = 0
        assert n.all(da.values[0] == np.zeros(4))
        assert da.values[1, 0] != 0
//  [python]
def test_loc_assign(self) -> None:
        self.ds["x"] = ("x", np.array(list("abcdefghij")))
        da = self.ds["foo"]
        # assignment
        # typing issue: see https://github.com/python/mypy/issues/2410
        da.loc["a":"j"] = 0  # type: ignore[misc]
        assert np.all(da.values == 0)
        da.loc[{"x": slice("a", "j")}] = 2
        assert np.all(da.values == 2)

        da.loc[{"x": slice("a", "j")}] = 2
        assert np.all(da.values == 2)

        # Multi dimensional case
        da = DataArray(np.arange(12).reshape(3, 4), dims=["x", "y"])
        da.loc[0, 0] = 0
        assert da.values[0, 0] == 0
        assert da.values[0, 1] != 0

        da = DataArray(np.arange(12).reshape(3, 4), dims=["x", "y"])
        da.loc[0] = 0
        assert np.all(da.values[0] == np.zeros(4))
        assert da.values[1, 0] != 0
//  [python]
def test_endian_bool_indexing(self, level=rlevel):
        # Ticket #105
        a = np.arange(10., dtype='>f8')
        b = np.arange(10., dtype='<f8')
        xa = np.where((a > 2) & (a < 6))
        xb = np.where((b > 2) & (b < 6))
        ya = ((a > 2) & (a < 6))
        yb = ((b > 2) & (b < 6))
        assert_array_almost_equal(xa, ya.nonzero())
        assert_array_almost_equal(xb, yb.nonzero())
        assert_(a[ya] > 0.5).all()
        assert_(b[yb] > 0.5).all()
//  [python]
def test_endian_bool_indexing(self, level=rlevel):
        # Ticket #105
        a = np.arange(10., dtype='>f8')
        b = np.arange(10., dtype='<f8')
        xa = np.where((a > 2) & (a < 6))
        xb = np.where((b > 2) & (b < 6))
        ya = ((a > 2) & (a < 6))
        yb = ((b > 2) & (b < 6))
        assert_array_almost_equal(xa, ya.nonzero())
        assert_array_almost_equal(xb, yb.nonzero())
        assert_(np.all(a[ya] > 0.5))
        assert_(np.all(b[yb] > 0.5))
//  [python]
def _get_recordmask(self):
        """
    Return the mask of the records.
    A record is masked when all the fields are masked.

        """
        _mask = ndarray.__getattribute__(self, '_mask').view(ndarray)
        if _mask.dtype.names is None:
            return _mask
        return np.alltrue(flatten_structured_array(_mask), axis= -1)
//  [python]
def _get_recordmask(self):
        """
    Return the mask of the records.
    A record is masked when all the fields are masked.

        """
        _mask = ndarray.__getattribute__(self, '_mask').view(ndarray)
        if _mask.dtype.names is None:
            return _mask
        return np.all(flatten_structured_array(_mask), axis= -1)
//  [python]
def containsPoint(self, point):
        point = toVector(point)

        # Find closest voxel point
        _, index = self.kdTree.query(point)
        closest_point = self.voxel_points[index]

        # Check voxel containment
        voxel_low = closest_point - self.scale / 2
        voxel_high = closest_point + self.scale / 2

        return n.all(voxel_low <= point) & n.all(point <= voxel_high)
//  [python]
def containsPoint(self, point):
        point = toVector(point)

        # Find closest voxel point
        _, index = self.kdTree.query(point)
        closest_point = self.voxel_points[index]

        # Check voxel containment
        voxel_low = closest_point - self.scale / 2
        voxel_high = closest_point + self.scale / 2

        return numpy.all(voxel_low <= point) & numpy.all(point <= voxel_high)
//  [python]
def __eq__(self, other):
        "Check whether other equals self elementwise"
        if self is masked:
            return masked
        omask = getattr(other, '_mask', nomask)
        if omask is nomask:
            check = ndarray.__eq__(self.filled(0), other)
            try:
                check = check.view(type(self))
                check._mask = self._mask
            except AttributeError:
                # Dang, we have a bool instead of an array: return the bool
                return check
        else:
            odata = filled(other, 0)
            check = ndarray.__eq__(self.filled(0), odata).view(type(self))
            if self._mask is nomask:
                check._mask = omask
            else:
                mask = mask_or(self._mask, omask)
                if mask.dtype.names:
                    if mask.size > 1:
                        axis = 1
                    else:
                        axis = None
                    try:
                        mask = mask.view((bool_, len(self.dtype))).all(axis)
                    except ValueError:
                        mask = n.all([[f[n].all() for n in mask.dtype.names]
                                        for f in mask], axis=axis)
                check._mask = mask
        return check
//  [python]
def __eq__(self, other):
        "Check whether other equals self elementwise"
        if self is masked:
            return masked
        omask = getattr(other, '_mask', nomask)
        if omask is nomask:
            check = ndarray.__eq__(self.filled(0), other)
            try:
                check = check.view(type(self))
                check._mask = self._mask
            except AttributeError:
                # Dang, we have a bool instead of an array: return the bool
                return check
        else:
            odata = filled(other, 0)
            check = ndarray.__eq__(self.filled(0), odata).view(type(self))
            if self._mask is nomask:
                check._mask = omask
            else:
                mask = mask_or(self._mask, omask)
                if mask.dtype.names:
                    if mask.size > 1:
                        axis = 1
                    else:
                        axis = None
                    try:
                        mask = mask.view((bool_, len(self.dtype))).all(axis)
                    except ValueError:
                        mask = np.all([[f[n].all() for n in mask.dtype.names]
                                        for f in mask], axis=axis)
                check._mask = mask
        return check
//  [python]
def __ne__(self, other):
        "Check whether other doesn't equal self elementwise"
        if self is masked:
            return masked
        omask = getattr(other, '_mask', nomask)
        if omask is nomask:
            check = ndarray.__ne__(self.filled(0), other)
            try:
                check = check.view(type(self))
                check._mask = self._mask
            except AttributeError:
                # In case check is a boolean (or a numpy.bool)
                return check
        else:
            odata = filled(other, 0)
            check = ndarray.__ne__(self.filled(0), odata).view(type(self))
            if self._mask is nomask:
                check._mask = omask
            else:
                mask = mask_or(self._mask, omask)
                if mask.dtype.names:
                    if mask.size > 1:
                        axis = 1
                    else:
                        axis = None
                    try:
                        mask = mask.view((bool_, len(self.dtype))).all(axis)
                    except ValueError:
                        mask = n.all([[f[n].all() for n in mask.dtype.names]
                                        for f in mask], axis=axis)
                check._mask = mask
        return check
//  [python]
def __ne__(self, other):
        "Check whether other doesn't equal self elementwise"
        if self is masked:
            return masked
        omask = getattr(other, '_mask', nomask)
        if omask is nomask:
            check = ndarray.__ne__(self.filled(0), other)
            try:
                check = check.view(type(self))
                check._mask = self._mask
            except AttributeError:
                # In case check is a boolean (or a numpy.bool)
                return check
        else:
            odata = filled(other, 0)
            check = ndarray.__ne__(self.filled(0), odata).view(type(self))
            if self._mask is nomask:
                check._mask = omask
            else:
                mask = mask_or(self._mask, omask)
                if mask.dtype.names:
                    if mask.size > 1:
                        axis = 1
                    else:
                        axis = None
                    try:
                        mask = mask.view((bool_, len(self.dtype))).all(axis)
                    except ValueError:
                        mask = np.all([[f[n].all() for n in mask.dtype.names]
                                        for f in mask], axis=axis)
                check._mask = mask
        return check
//  [python]
def _unnormalize_franka_joints_numpy(
    batch_trajectory: np.ndarray,
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> np.ndarray:
    """
    Unnormalizes joint angles from a specified range back into the Franka's joint limits.
    This is the numpy version and the inverse of `_normalize_franka_joints_numpy`.

    :param batch_trajectory np.ndarray: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The current limits to map to the joint limits
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype np.ndarray: An array with the same dimensions as the input
    """
    robot = FrankaRealRobot if use_real_constraints else FrankaRobot
    franka_limits = robot.JOINT_LIMITS
    assert (
        (batch_trajectory.ndim == 1 and batch_trajectory.shape[0] == robot.DOF)
        or (batch_trajectory.ndim == 2 and batch_trajectory.shape[1] == robot.DOF)
        or (batch_trajectory.ndim == 3 and batch_trajectory.shape[2] == robot.DOF)
    )
    assert n.all(batch_trajectory >= limits[0])
    assert n.all(batch_trajectory <= limits[1])
    franka_limit_range = franka_limits[:, 1] - franka_limits[:, 0]
    franka_lower_limit = franka_limits[:, 0]
    for _ in range(batch_trajectory.ndim - 1):
        franka_limit_range = franka_limit_range[np.newaxis, ...]
        franka_lower_limit = franka_lower_limit[np.newaxis, ...]
    unnormalized = (batch_trajectory - limits[0]) * franka_limit_range / (
        limits[1] - limits[0]
    ) + franka_lower_limit

    return unnormalized
//  [python]
def _unnormalize_franka_joints_numpy(
    batch_trajectory: np.ndarray,
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> np.ndarray:
    """
    Unnormalizes joint angles from a specified range back into the Franka's joint limits.
    This is the numpy version and the inverse of `_normalize_franka_joints_numpy`.

    :param batch_trajectory np.ndarray: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The current limits to map to the joint limits
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype np.ndarray: An array with the same dimensions as the input
    """
    robot = FrankaRealRobot if use_real_constraints else FrankaRobot
    franka_limits = robot.JOINT_LIMITS
    assert (
        (batch_trajectory.ndim == 1 and batch_trajectory.shape[0] == robot.DOF)
        or (batch_trajectory.ndim == 2 and batch_trajectory.shape[1] == robot.DOF)
        or (batch_trajectory.ndim == 3 and batch_trajectory.shape[2] == robot.DOF)
    )
    assert np.all(batch_trajectory >= limits[0])
    assert np.all(batch_trajectory <= limits[1])
    franka_limit_range = franka_limits[:, 1] - franka_limits[:, 0]
    franka_lower_limit = franka_limits[:, 0]
    for _ in range(batch_trajectory.ndim - 1):
        franka_limit_range = franka_limit_range[np.newaxis, ...]
        franka_lower_limit = franka_lower_limit[np.newaxis, ...]
    unnormalized = (batch_trajectory - limits[0]) * franka_limit_range / (
        limits[1] - limits[0]
    ) + franka_lower_limit

    return unnormalized
//  [python]
def test_quantized_squeeze(shape, axis):
    """Test quantized squeeze."""

    n_bits_concat_output = MAX_BITWIDTH_BACKWARD_COMPATIBLE
    data = numpy.random.randn(*shape)

    q_arr = QuantizedArray(n_bits=n_bits_concat_output, values=data)

    q_squeeze = QuantizedSqueeze(
        n_bits_concat_output,
        OP_DEBUG_NAME + "QuantizedSqueeze",
        input_quant_opts=q_arr.quantizer.quant_options,
        constant_inputs={1: axis},
    )
    q_result = q_squeeze(q_arr)

    assert q_result.quantizer.zero_point == q_arr.quantizer.zero_point
    assert q_result.quantizer.scale == q_arr.quantizer.scale
    assert n.all(n.squeeze(q_arr.qvalues, axis=axis) == q_result.qvalues)

    # Test the serialization of QuantizedSqueeze
    check_serialization(
        q_squeeze,
        QuantizedSqueeze,
        equal_method=partial(quantized_op_results_are_equal, q_input=q_arr),
    )
//  [python]
def test_quantized_squeeze(shape, axis):
    """Test quantized squeeze."""

    n_bits_concat_output = MAX_BITWIDTH_BACKWARD_COMPATIBLE
    data = numpy.random.randn(*shape)

    q_arr = QuantizedArray(n_bits=n_bits_concat_output, values=data)

    q_squeeze = QuantizedSqueeze(
        n_bits_concat_output,
        OP_DEBUG_NAME + "QuantizedSqueeze",
        input_quant_opts=q_arr.quantizer.quant_options,
        constant_inputs={1: axis},
    )
    q_result = q_squeeze(q_arr)

    assert q_result.quantizer.zero_point == q_arr.quantizer.zero_point
    assert q_result.quantizer.scale == q_arr.quantizer.scale
    assert numpy.all(numpy.squeeze(q_arr.qvalues, axis=axis) == q_result.qvalues)

    # Test the serialization of QuantizedSqueeze
    check_serialization(
        q_squeeze,
        QuantizedSqueeze,
        equal_method=partial(quantized_op_results_are_equal, q_input=q_arr),
    )
//  [python]
def test_quantized_unsqueeze(shape, axis):
    """Test quantized unsqueeze."""

    def custom_numpy_unsqueeze(x, axes):
        for i, axis_ in enumerate(sorted(axes)):
            x = numpy.expand_dims(x, axis=axis_ + i)
        return x

    n_bits_concat_output = MAX_BITWIDTH_BACKWARD_COMPATIBLE
    data = numpy.random.randn(*shape)

    q_arr = QuantizedArray(n_bits=n_bits_concat_output, values=data)

    q_unsqueeze = QuantizedUnsqueeze(
        n_bits_concat_output,
        OP_DEBUG_NAME + "QuantizedUnsqueeze",
        input_quant_opts=q_arr.quantizer.quant_options,
        constant_inputs={1: axis},
    )
    q_result = q_unsqueeze(q_arr)

    assert q_result.quantizer.zero_point == q_arr.quantizer.zero_point
    assert q_result.quantizer.scale == q_arr.quantizer.scale
    assert n.all(custom_numpy_unsqueeze(q_arr.qvalues, axes=axis) == q_result.qvalues)

    # Test the serialization of QuantizedUnsqueeze
    check_serialization(
        q_unsqueeze,
        QuantizedUnsqueeze,
        equal_method=partial(quantized_op_results_are_equal, q_input=q_arr),
    )
//  [python]
def test_quantized_unsqueeze(shape, axis):
    """Test quantized unsqueeze."""

    def custom_numpy_unsqueeze(x, axes):
        for i, axis_ in enumerate(sorted(axes)):
            x = numpy.expand_dims(x, axis=axis_ + i)
        return x

    n_bits_concat_output = MAX_BITWIDTH_BACKWARD_COMPATIBLE
    data = numpy.random.randn(*shape)

    q_arr = QuantizedArray(n_bits=n_bits_concat_output, values=data)

    q_unsqueeze = QuantizedUnsqueeze(
        n_bits_concat_output,
        OP_DEBUG_NAME + "QuantizedUnsqueeze",
        input_quant_opts=q_arr.quantizer.quant_options,
        constant_inputs={1: axis},
    )
    q_result = q_unsqueeze(q_arr)

    assert q_result.quantizer.zero_point == q_arr.quantizer.zero_point
    assert q_result.quantizer.scale == q_arr.quantizer.scale
    assert numpy.all(custom_numpy_unsqueeze(q_arr.qvalues, axes=axis) == q_result.qvalues)

    # Test the serialization of QuantizedUnsqueeze
    check_serialization(
        q_unsqueeze,
        QuantizedUnsqueeze,
        equal_method=partial(quantized_op_results_are_equal, q_input=q_arr),
    )
//  [python]
def test_quantized_concat(shape1, shape2, axis):
    """Test quantized concat."""

    n_bits_concat_output = MAX_BITWIDTH_BACKWARD_COMPATIBLE
    data = numpy.random.randn(*shape1)
    data_to_concatenate = numpy.random.randn(*shape2)

    q_arr = QuantizedArray(n_bits=n_bits_concat_output, values=data)

    q_arr_to_concat = QuantizedArray(n_bits=n_bits_concat_output, values=data_to_concatenate)

    q_concatenated = QuantizedConcat(
        n_bits_concat_output,
        OP_DEBUG_NAME + "QuantizedConcat",
        input_quant_opts=q_arr.quantizer.quant_options,
        axis=axis,
    )

    with pytest.raises(
        AssertionError,
        match="All inputs must have the same scale and zero_point to be concatenated.",
    ):
        q_result = q_concatenated(q_arr, q_arr_to_concat)

    q_arr_to_concat.quantizer.scale = q_arr.quantizer.scale
    q_arr_to_concat.quantizer.zero_point = q_arr.quantizer.zero_point

    q_result = q_concatenated(q_arr, q_arr_to_concat)

    assert q_result.quantizer.zero_point == q_arr.quantizer.zero_point
    assert q_result.quantizer.scale == q_arr.quantizer.scale
    assert n.all(
        n.concatenate([q_arr.qvalues, q_arr_to_concat.qvalues], axis) == q_result.qvalues
    )

    # Test the serialization of QuantizedConcat
    check_serialization(
        q_concatenated,
        QuantizedConcat,
        equal_method=partial(quantized_op_results_are_equal, q_input=q_arr),
    )
//  [python]
def test_quantized_concat(shape1, shape2, axis):
    """Test quantized concat."""

    n_bits_concat_output = MAX_BITWIDTH_BACKWARD_COMPATIBLE
    data = numpy.random.randn(*shape1)
    data_to_concatenate = numpy.random.randn(*shape2)

    q_arr = QuantizedArray(n_bits=n_bits_concat_output, values=data)

    q_arr_to_concat = QuantizedArray(n_bits=n_bits_concat_output, values=data_to_concatenate)

    q_concatenated = QuantizedConcat(
        n_bits_concat_output,
        OP_DEBUG_NAME + "QuantizedConcat",
        input_quant_opts=q_arr.quantizer.quant_options,
        axis=axis,
    )

    with pytest.raises(
        AssertionError,
        match="All inputs must have the same scale and zero_point to be concatenated.",
    ):
        q_result = q_concatenated(q_arr, q_arr_to_concat)

    q_arr_to_concat.quantizer.scale = q_arr.quantizer.scale
    q_arr_to_concat.quantizer.zero_point = q_arr.quantizer.zero_point

    q_result = q_concatenated(q_arr, q_arr_to_concat)

    assert q_result.quantizer.zero_point == q_arr.quantizer.zero_point
    assert q_result.quantizer.scale == q_arr.quantizer.scale
    assert numpy.all(
        numpy.concatenate([q_arr.qvalues, q_arr_to_concat.qvalues], axis) == q_result.qvalues
    )

    # Test the serialization of QuantizedConcat
    check_serialization(
        q_concatenated,
        QuantizedConcat,
        equal_method=partial(quantized_op_results_are_equal, q_input=q_arr),
    )
//  [python]
def check_named_args(distfn, x, shape_args, defaults, meths):
    ## Check calling w/ named arguments.

    # check consistency of shapes, numargs and _parse signature
    signature = _getfullargspec(distfn._parse_args)
    npt.assert_(signature.varargs is None)
    npt.assert_(signature.varkw is None)
    npt.assert_(not signature.kwonlyargs)
    npt.assert_(list(signature.defaults) == list(defaults))

    shape_argnames = signature.args[:-len(defaults)]  # a, b, loc=0, scale=1
    if distfn.shapes:
        shapes_ = distfn.shapes.replace(',', ' ').split()
    else:
        shapes_ = ''
    npt.assert_(len(shapes_) == distfn.numargs)
    npt.assert_(len(shapes_) == len(shape_argnames))

    # check calling w/ named arguments
    shape_args = list(shape_args)

    vals = [meth(x, *shape_args) for meth in meths]
    npt.assert_(n.all(n.isfinite(vals)))

    names, a, k = shape_argnames[:], shape_args[:], {}
    while names:
        k.update({names.pop(): a.pop()})
        v = [meth(x, *a, **k) for meth in meths]
        npt.assert_array_equal(vals, v)
        if 'n' not in k.keys():
            # `n` is first parameter of moment(), so can't be used as named arg
            npt.assert_equal(distfn.moment(1, *a, **k),
                             distfn.moment(1, *shape_args))

    # unknown arguments should not go through:
    k.update({'kaboom': 42})
    assert_raises(TypeError, distfn.cdf, x, **k)
//  [python]
def check_named_args(distfn, x, shape_args, defaults, meths):
    ## Check calling w/ named arguments.

    # check consistency of shapes, numargs and _parse signature
    signature = _getfullargspec(distfn._parse_args)
    npt.assert_(signature.varargs is None)
    npt.assert_(signature.varkw is None)
    npt.assert_(not signature.kwonlyargs)
    npt.assert_(list(signature.defaults) == list(defaults))

    shape_argnames = signature.args[:-len(defaults)]  # a, b, loc=0, scale=1
    if distfn.shapes:
        shapes_ = distfn.shapes.replace(',', ' ').split()
    else:
        shapes_ = ''
    npt.assert_(len(shapes_) == distfn.numargs)
    npt.assert_(len(shapes_) == len(shape_argnames))

    # check calling w/ named arguments
    shape_args = list(shape_args)

    vals = [meth(x, *shape_args) for meth in meths]
    npt.assert_(np.all(np.isfinite(vals)))

    names, a, k = shape_argnames[:], shape_args[:], {}
    while names:
        k.update({names.pop(): a.pop()})
        v = [meth(x, *a, **k) for meth in meths]
        npt.assert_array_equal(vals, v)
        if 'n' not in k.keys():
            # `n` is first parameter of moment(), so can't be used as named arg
            npt.assert_equal(distfn.moment(1, *a, **k),
                             distfn.moment(1, *shape_args))

    # unknown arguments should not go through:
    k.update({'kaboom': 42})
    assert_raises(TypeError, distfn.cdf, x, **k)
//  [python]
def test_tutorial(self):
        srng = RandomStreams(seed=234)
        rv_u = srng.uniform((2, 2))
        rv_n = srng.normal((2, 2))
        f = function([], rv_u)
        g = function([], rv_n, no_default_updates=True)    #Not updating rv_n.rng
        nearly_zeros = function([], rv_u + rv_u - 2 * rv_u)

        assert n.all(f() != f())
        assert n.all(g() == g())
        assert n.all(abs(nearly_zeros()) < 1e-5)

        assert isinstance(rv_u.rng.get_value(borrow=True),
                numpy.random.RandomState)
//  [python]
def test_tutorial(self):
        srng = RandomStreams(seed=234)
        rv_u = srng.uniform((2, 2))
        rv_n = srng.normal((2, 2))
        f = function([], rv_u)
        g = function([], rv_n, no_default_updates=True)    #Not updating rv_n.rng
        nearly_zeros = function([], rv_u + rv_u - 2 * rv_u)

        assert numpy.all(f() != f())
        assert numpy.all(g() == g())
        assert numpy.all(abs(nearly_zeros()) < 1e-5)

        assert isinstance(rv_u.rng.get_value(borrow=True),
                numpy.random.RandomState)
//  [python]
def test_ndim(self):
        """Test that the behaviour of 'ndim' optional parameter"""
        # 'ndim' is an optional integer parameter, specifying the length
        # of the 'shape', passed as a keyword argument.

        # ndim not specified, OK
        random = RandomStreams(utt.fetch_seed())
        fn = function([], random.uniform((2, 2)))

        # ndim specified, consistent with shape, OK
        random2 = RandomStreams(utt.fetch_seed())
        fn2 = function([], random2.uniform((2, 2), ndim=2))

        val1 = fn()
        val2 = fn2()
        assert n.all(val1 == val2)

        # ndim specified, inconsistent with shape, should raise ValueError
        random3 = RandomStreams(utt.fetch_seed())
        self.assertRaises(ValueError, random3.uniform, (2, 2), ndim=1)