MCPcopy Create free account
hub / github.com/numpy/numpy / npyiter_coalesce_axes

Function npyiter_coalesce_axes

numpy/core/src/multiarray/nditer_api.c:1679–1746  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1677}
1678
1679NPY_NO_EXPORT void
1680npyiter_coalesce_axes(NpyIter *iter)
1681{
1682 npy_uint32 itflags = NIT_ITFLAGS(iter);
1683 int idim, ndim = NIT_NDIM(iter);
1684 int nop = NIT_NOP(iter);
1685
1686 npy_intp istrides, nstrides = NAD_NSTRIDES();
1687 NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
1688 npy_intp sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, nop);
1689 NpyIter_AxisData *ad_compress = axisdata;
1690 npy_intp new_ndim = 1;
1691
1692 /* The HASMULTIINDEX or IDENTPERM flags do not apply after coalescing */
1693 NIT_ITFLAGS(iter) &= ~(NPY_ITFLAG_IDENTPERM|NPY_ITFLAG_HASMULTIINDEX);
1694
1695 for (idim = 0; idim < ndim-1; ++idim) {
1696 int can_coalesce = 1;
1697 npy_intp shape0 = NAD_SHAPE(ad_compress);
1698 npy_intp shape1 = NAD_SHAPE(NIT_INDEX_AXISDATA(axisdata, 1));
1699 npy_intp *strides0 = NAD_STRIDES(ad_compress);
1700 npy_intp *strides1 = NAD_STRIDES(NIT_INDEX_AXISDATA(axisdata, 1));
1701
1702 /* Check that all the axes can be coalesced */
1703 for (istrides = 0; istrides < nstrides; ++istrides) {
1704 if (!((shape0 == 1 && strides0[istrides] == 0) ||
1705 (shape1 == 1 && strides1[istrides] == 0)) &&
1706 (strides0[istrides]*shape0 != strides1[istrides])) {
1707 can_coalesce = 0;
1708 break;
1709 }
1710 }
1711
1712 if (can_coalesce) {
1713 npy_intp *strides = NAD_STRIDES(ad_compress);
1714
1715 NIT_ADVANCE_AXISDATA(axisdata, 1);
1716 NAD_SHAPE(ad_compress) *= NAD_SHAPE(axisdata);
1717 for (istrides = 0; istrides < nstrides; ++istrides) {
1718 if (strides[istrides] == 0) {
1719 strides[istrides] = NAD_STRIDES(axisdata)[istrides];
1720 }
1721 }
1722 }
1723 else {
1724 NIT_ADVANCE_AXISDATA(axisdata, 1);
1725 NIT_ADVANCE_AXISDATA(ad_compress, 1);
1726 if (ad_compress != axisdata) {
1727 memcpy(ad_compress, axisdata, sizeof_axisdata);
1728 }
1729 ++new_ndim;
1730 }
1731 }
1732
1733 /*
1734 * If the number of axes shrunk, reset the perm and
1735 * compress the data into the new layout.
1736 */

Callers 2

NpyIter_AdvancedNewFunction · 0.85
NpyIter_RemoveMultiIndexFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected