jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __voxel_4conectedness(shape):
    """
    Returns the number of edges for the supplied image shape assuming 4-connectedness.

    The name of the function has historical reasons. Essentially it returns the number
    of edges assuming 4-connectedness only for 2D. For 3D it assumes 6-connectedness,
    etc.

    @param shape the shape of the image
    @type shape sequence
    @return the number of edges
    @rtype int
    """
    shape = list(shape)
    while 1 in shape:
        shape.remove(
            1
        )  # empty resp. 1-sized dimensions have to be removed (equal to numpy.squeeze on the array)
    return int(
        round(sum([(dim - 1) / float(dim) for dim in shape]) * numpy.prod(shape))
    )
//  [python]
def example_dataset(shape):
    # TODO: simplify the creation of datasets here
    # TODO: See https://github.com/pangeo-data/rechunker/pull/93#discussion_r713939185
    # TODO: Maybe it is best to refactor tests to use `chunk_ds`
    xarray = pytest.importorskip("xarray")

    a = numpy.arange(numpy.prod(shape)).reshape(shape).astype("f4")
    a[-1] = numpy.nan
    ds = xarray.Dataset(
        dict(
            a=xarray.DataArray(
                a, dims=["x", "y"], attrs={"a1": 1, "a2": [1, 2, 3], "a3": "x"}
            ),
            b=xarray.DataArray(numpy.ones(shape[0]), dims=["x"]),
            c=xarray.DataArray(numpy.ones(shape[1]), dims=["y"]),
        ),
        coords=dict(
            cx=xarray.DataArray(numpy.ones(shape[0]), dims=["x"]),
            cy=xarray.DataArray(numpy.ones(shape[1]), dims=["y"]),
        ),
        attrs={"a1": 1, "a2": [1, 2, 3], "a3": "x"},
    )
    return ds
//  [python]
def setUp(self):
        super(PointSelectionTestCase, self).setUp()
        # Limits for selections
        self.limits = [
            (0, 1),  # just one element
            (20, -10),  # no elements
            (-10, 4),  # several elements
            (0, 10),   # several elements (again)
        ]

        # Create a sample array
        size = n.prod(self.shape)
        nparr = n.arange(size, dtype=numpy.int32).reshape(self.shape)
        self.nparr = nparr
        self.tbarr = self.h5file.create_array(self.h5file.root, 'array', nparr)
//  [python]
def setUp(self):
        super(PointSelectionTestCase, self).setUp()
        # Limits for selections
        self.limits = [
            (0, 1),  # just one element
            (20, -10),  # no elements
            (-10, 4),  # several elements
            (0, 10),   # several elements (again)
        ]

        # Create a sample array
        size = numpy.prod(self.shape)
        nparr = numpy.arange(size, dtype=numpy.int32).reshape(self.shape)
        self.nparr = nparr
        self.tbarr = self.h5file.create_array(self.h5file.root, 'array', nparr)
//  [python]
def test(self):
        """Test correct broadcasting when the array atom is not scalar."""

        array_shape = (2, 3)
        element_shape = (3,)

        dtype = numpy.dtype((numpy.int, element_shape))
        atom = Atom.from_dtype(dtype)
        h5arr = self.h5file.create_carray(self.h5file.root, 'array',
                                          atom, array_shape)

        size = n.prod(element_shape)
        nparr = numpy.arange(size).reshape(element_shape)

        h5arr[0] = nparr
        self.assertTrue(numpy.all(h5arr[0] == nparr))
//  [python]
def test(self):
        """Test correct broadcasting when the array atom is not scalar."""

        array_shape = (2, 3)
        element_shape = (3,)

        dtype = numpy.dtype((numpy.int, element_shape))
        atom = Atom.from_dtype(dtype)
        h5arr = self.h5file.create_carray(self.h5file.root, 'array',
                                          atom, array_shape)

        size = numpy.prod(element_shape)
        nparr = numpy.arange(size).reshape(element_shape)

        h5arr[0] = nparr
        self.assertTrue(numpy.all(h5arr[0] == nparr))
//  [python]
def _mom(self, k, mean, sigma, cache):
        if isinstance(mean, chaospy.Distribution):
            mean = mean._get_cache(None, cache=cache, get=0)
            if isinstance(mean, chaospy.Distribution):
                raise chaospy.UnsupportedFeature(
                    "Analytical moment of a conditional not supported"
                )
        out = 0.0
        for idx, kdx in enumerate(numpy.ndindex(*[_ + 1 for _ in k])):
            coef = n.prod(special.comb(k.T, kdx).T, 0)
            diff = k.T - kdx
            pos = diff >= 0
            diff = diff * pos
            pos = numpy.all(pos)
            location_ = n.prod(mean**diff, -1)

            out = out + pos * coef * location_ * isserlis_moment(tuple(kdx), sigma)

        return out
//  [python]
def _mom(self, k, mean, sigma, cache):
        if isinstance(mean, chaospy.Distribution):
            mean = mean._get_cache(None, cache=cache, get=0)
            if isinstance(mean, chaospy.Distribution):
                raise chaospy.UnsupportedFeature(
                    "Analytical moment of a conditional not supported"
                )
        out = 0.0
        for idx, kdx in enumerate(numpy.ndindex(*[_ + 1 for _ in k])):
            coef = numpy.prod(special.comb(k.T, kdx).T, 0)
            diff = k.T - kdx
            pos = diff >= 0
            diff = diff * pos
            pos = numpy.all(pos)
            location_ = numpy.prod(mean**diff, -1)

            out = out + pos * coef * location_ * isserlis_moment(tuple(kdx), sigma)

        return out
//  [python]
def fake_Cartesian(Nd):
    dim = len(Nd) # dimension
    M = n.prod(Nd)
    om = numpy.zeros((M, dim), dtype = numpy.float64)
    grid = numpy.indices(Nd)
    for dimid in range(0, dim):
        om[:, dimid] = (grid[dimid].ravel() *2/ Nd[dimid] - 1.0)*numpy.pi
    return om
//  [python]
def fake_Cartesian(Nd):
    dim = len(Nd) # dimension
    M = numpy.prod(Nd)
    om = numpy.zeros((M, dim), dtype = numpy.float64)
    grid = numpy.indices(Nd)
    for dimid in range(0, dim):
        om[:, dimid] = (grid[dimid].ravel() *2/ Nd[dimid] - 1.0)*numpy.pi
    return om
//  [python]
def DFT_matrix(Nd, om=None):
    dim = len(Nd) # dimension
    if om is None:
        om = fake_Cartesian(Nd)
    N = n.prod(Nd)
    omN = numpy.zeros((N, dim), dtype = numpy.float64)
    grid = numpy.indices(Nd)    
    for dimid in range(0, dim):
        omN[:, dimid] = (grid[dimid].ravel() - Nd[dimid]/2 )
    M = om.shape[0]
    A = numpy.einsum('m, n -> mn', om[:, 0], omN[:,0], optimize='optimal')
    for d in range(1, dim):
        A += numpy.einsum('m, n -> mn', om[:, d], omN[:,d], optimize='optimal')
        
    return numpy.exp(-1.0j* A)
//  [python]
def DFT_matrix(Nd, om=None):
    dim = len(Nd) # dimension
    if om is None:
        om = fake_Cartesian(Nd)
    N = numpy.prod(Nd)
    omN = numpy.zeros((N, dim), dtype = numpy.float64)
    grid = numpy.indices(Nd)    
    for dimid in range(0, dim):
        omN[:, dimid] = (grid[dimid].ravel() - Nd[dimid]/2 )
    M = om.shape[0]
    A = numpy.einsum('m, n -> mn', om[:, 0], omN[:,0], optimize='optimal')
    for d in range(1, dim):
        A += numpy.einsum('m, n -> mn', om[:, d], omN[:,d], optimize='optimal')
        
    return numpy.exp(-1.0j* A)
//  [python]
def test_neibs_manual_step(self):
        shape = (2, 3, 5, 5)
        for dtype in self.dtypes:
            images = shared(numpy.asarray(numpy.arange(numpy.prod(
                            shape)).reshape(shape), dtype=dtype))
            neib_shape = T.as_tensor_variable((3, 3))
            neib_step = T.as_tensor_variable((2, 2))
            for border in ['valid', 'ignore_borders']:
                f = function([],
                             images2neibs(images, neib_shape, neib_step,
                                          mode=border),
                             mode=self.mode)

                neibs = f()
                assert self.op in [type(node.op)
                                   for node in f.maker.fgraph.toposort()]

                assert numpy.allclose(neibs,
              [[  0,   1,   2,   5,   6,   7,  10,  11,  12],
               [  2,   3,   4,   7,   8,   9,  12,  13,  14],
               [ 10,  11,  12,  15,  16,  17,  20,  21,  22],
               [ 12,  13,  14,  17,  18,  19,  22,  23,  24],
               [ 25,  26,  27,  30,  31,  32,  35,  36,  37],
               [ 27,  28,  29,  32,  33,  34,  37,  38,  39],
               [ 35,  36,  37,  40,  41,  42,  45,  46,  47],
               [ 37,  38,  39,  42,  43,  44,  47,  48,  49],
               [ 50,  51,  52,  55,  56,  57,  60,  61,  62],
               [ 52,  53,  54,  57,  58,  59,  62,  63,  64],
               [ 60,  61,  62,  65,  66,  67,  70,  71,  72],
               [ 62,  63,  64,  67,  68,  69,  72,  73,  74],
               [ 75,  76,  77,  80,  81,  82,  85,  86,  87],
               [ 77,  78,  79,  82,  83,  84,  87,  88,  89],
               [ 85,  86,  87,  90,  91,  92,  95,  96,  97],
               [ 87,  88,  89,  92,  93,  94,  97,  98,  99],
               [100, 101, 102, 105, 106, 107, 110, 111, 112],
               [102, 103, 104, 107, 108, 109, 112, 113, 114],
               [110, 111, 112, 115, 116, 117, 120, 121, 122],
               [112, 113, 114, 117, 118, 119, 122, 123, 124],
               [125, 126, 127, 130, 131, 132, 135, 136, 137],
               [127, 128, 129, 132, 133, 134, 137, 138, 139],
               [135, 136, 137, 140, 141, 142, 145, 146, 147],
               [137, 138, 139, 142, 143, 144, 147, 148, 149]])
//  [python]
def plan(self, om, Nd):
#         if batch != None:
#             self.batch = numpy.int(batch)
        self.Nd = Nd
        self.ndims = len(Nd)
        self.F_matrix = []
        compute_str_x = ''
        F_str = ''
        for dimid in range(0, self.ndims):
            self.F_matrix += [DFT_matrix((Nd[dimid],), om[:, dimid:dimid+1]),]
            compute_str_x += lower_case[dimid]
            F_str += ', m' + lower_case[dimid]
        
#         if type(batch)==numpy.int:
#             self.compute_str_forward = compute_str_x +'z' + F_str + '-> mz'
#             self.compute_str_adj = 'm' + F_str + '->' + compute_str_x +'z'
#         else: 
        self.compute_str_forward = compute_str_x + F_str + '-> m'
        self.compute_str_adj = 'm' + F_str + '->' + compute_str_x
        
        self.scale = n.prod(Nd)
//  [python]
def plan(self, om, Nd):
#         if batch != None:
#             self.batch = numpy.int(batch)
        self.Nd = Nd
        self.ndims = len(Nd)
        self.F_matrix = []
        compute_str_x = ''
        F_str = ''
        for dimid in range(0, self.ndims):
            self.F_matrix += [DFT_matrix((Nd[dimid],), om[:, dimid:dimid+1]),]
            compute_str_x += lower_case[dimid]
            F_str += ', m' + lower_case[dimid]
        
#         if type(batch)==numpy.int:
#             self.compute_str_forward = compute_str_x +'z' + F_str + '-> mz'
#             self.compute_str_adj = 'm' + F_str + '->' + compute_str_x +'z'
#         else: 
        self.compute_str_forward = compute_str_x + F_str + '-> m'
        self.compute_str_adj = 'm' + F_str + '->' + compute_str_x
        
        self.scale = numpy.prod(Nd)
//  [python]
def __call__(self):
    shape_data = self.f.read(self.shape_dtype.itemsize * self.dims)
    shape = numpy.fromstring(shape_data, self.shape_dtype)
    if self.byteswap:
      shape = shape.byteswap()
    
    len_data = int(numpy.multiply.reduce(shape)) * self.store_dtype.itemsize
    data = self.f.read(len_data)
    if len(data)!=len_data:
      raise EOFError('Ran out data reading array.')
    
    ret = numpy.fromstring(data, self.store_dtype).reshape(shape)
    if self.byteswap_data:
      ret = ret.byteswap()
      
    return ret
//  [python]
def __call__(self):
    shape_data = self.f.read(self.shape_dtype.itemsize * self.dims)
    shape = numpy.fromstring(shape_data, self.shape_dtype)
    if self.byteswap:
      shape = shape.byteswap()
    
    len_data = int(numpy.prod(shape)) * self.store_dtype.itemsize
    data = self.f.read(len_data)
    if len(data)!=len_data:
      raise EOFError('Ran out data reading array.')
    
    ret = numpy.fromstring(data, self.store_dtype).reshape(shape)
    if self.byteswap_data:
      ret = ret.byteswap()
      
    return ret
//  [python]
def size(self):
        """The total number of elements in the represented signal."""
        return np.prod(self.shape)
//  [python]
def _isDissipative(obj):
    """
    Determine whether this combination of potentials and forces is Dissipative

    Parameters
    ----------
    obj : Potential/DissipativeForce instance or list of such instances

    Returns
    -------
    bool
        True or False depending on whether the object is dissipative

    Notes
    -----
    - 2018-03-16 - Written - Bovy (UofT)
    - 2023-05-29 - Adjusted to also take planarDissipativeForces into account - Bovy (UofT)

    """
    from .planarDissipativeForce import planarDissipativeForce
    from .Potential import flatten

    obj = flatten(obj)
    isList = isinstance(obj, list)
    if isList:
        isCons = [
            not isinstance(p, DissipativeForce)
            and not isinstance(p, planarDissipativeForce)
            for p in obj
        ]
        nonCons = not n.prod(n.array(isCons))
    else:
        nonCons = isinstance(obj, DissipativeForce) or isinstance(
            obj, planarDissipativeForce
        )
    return nonCons
//  [python]
def _isDissipative(obj):
    """
    Determine whether this combination of potentials and forces is Dissipative

    Parameters
    ----------
    obj : Potential/DissipativeForce instance or list of such instances

    Returns
    -------
    bool
        True or False depending on whether the object is dissipative

    Notes
    -----
    - 2018-03-16 - Written - Bovy (UofT)
    - 2023-05-29 - Adjusted to also take planarDissipativeForces into account - Bovy (UofT)

    """
    from .planarDissipativeForce import planarDissipativeForce
    from .Potential import flatten

    obj = flatten(obj)
    isList = isinstance(obj, list)
    if isList:
        isCons = [
            not isinstance(p, DissipativeForce)
            and not isinstance(p, planarDissipativeForce)
            for p in obj
        ]
        nonCons = not numpy.prod(numpy.array(isCons))
    else:
        nonCons = isinstance(obj, DissipativeForce) or isinstance(
            obj, planarDissipativeForce
        )
    return nonCons
//  [python]
def sample(self, rng=np.random):
        if isinstance(self.init, Distribution):
            # we sample this way so that any variancescaling distribution based
            # on n/d is scaled appropriately
            kernel = [
                self.init.sample(
                    self.input_shape.n_channels // self.groups, self.n_filters, rng=rng
                )
                for _ in range(n.prod(self.kernel_size))
            ]
            kernel = np.reshape(kernel, self.kernel_shape)
        else:
            kernel = np.array(self.init, dtype=rc.float_dtype)
        return kernel
//  [python]
def sample(self, rng=np.random):
        if isinstance(self.init, Distribution):
            # we sample this way so that any variancescaling distribution based
            # on n/d is scaled appropriately
            kernel = [
                self.init.sample(
                    self.input_shape.n_channels // self.groups, self.n_filters, rng=rng
                )
                for _ in range(np.prod(self.kernel_size))
            ]
            kernel = np.reshape(kernel, self.kernel_shape)
        else:
            kernel = np.array(self.init, dtype=rc.float_dtype)
        return kernel
//  [python]
def forward(self, x):
        """
        Args:
            x (CTensor): matrix.
        Return:
            the result CTensor
        """
        self._shape = x.shape()
        shape = list(self.shape)
        # handle the shape with 0
        shape = [
            self._shape[i]
            if i < len(self._shape) and shape[i] == 0 else shape[i]
            for i in range(len(shape))
        ]
        # handle the shape with -1
        hidden_shape = int(n.prod(self._shape) // n.abs(n.prod(shape)))
        self.cache = [int(s) if s != -1 else hidden_shape for s in shape]
        return singa.Reshape(x, self.cache)
//  [python]
def forward(self, x):
        """
        Args:
            x (CTensor): matrix.
        Return:
            the result CTensor
        """
        self._shape = x.shape()
        shape = list(self.shape)
        # handle the shape with 0
        shape = [
            self._shape[i]
            if i < len(self._shape) and shape[i] == 0 else shape[i]
            for i in range(len(shape))
        ]
        # handle the shape with -1
        hidden_shape = int(np.prod(self._shape) // np.abs(np.prod(shape)))
        self.cache = [int(s) if s != -1 else hidden_shape for s in shape]
        return singa.Reshape(x, self.cache)
//  [python]
def approximate_moment(
    distribution, k_loc, order=100000, rule="clenshaw_curtis", **kwargs
):
    """
    Approximation method for estimation of raw statistical moments.

    Uses quadrature integration to estimate the values.

    Args:
        distribution (Distribution):
            Distribution domain with dim=len(distribution)
        k_loc (Sequence[int, ...]):
            The exponents of the moments of interest with ``shape == (dim,)``.
        order (int):
            The quadrature order used in approximation.
        rule (str):
            Quadrature rule for integrating moments.
        kwargs:
            Extra args passed to :func:`chaospy.generate_quadrature`.

    Examples:
        >>> distribution = chaospy.Uniform(1, 4)
        >>> round(chaospy.approximate_moment(distribution, (1,)), 4)
        2.5
        >>> round(chaospy.approximate_moment(distribution, (2,)), 4)
        7.0

    """
    assert isinstance(distribution, chaospy.Distribution)
    k_loc = numpy.asarray(k_loc)
    if len(distribution) > 1:
        assert (
            not distribution.stochastic_dependent
        ), "Dependent distributions does not support moment approximation."
        assert len(k_loc) == len(distribution), "incorrect size of exponents"
        return n.prod(
            [
                approximate_moment(
                    distribution[idx], (k_loc[idx],), order=order, rule=rule, **kwargs
                )
                for idx in range(len(distribution))
            ],
            axis=0,
        )

    k_loc = tuple(k_loc.tolist())
    assert all(
        [isinstance(k, int) for k in k_loc]
    ), "exponents must be integers: %s found" % type(k_loc[0])

    order = int(1e5 if order is None else order)
    if (distribution, order) not in MOMENTS_QUADS:
        MOMENTS_QUADS[distribution, order] = chaospy.generate_quadrature(
            order, distribution, rule=rule, **kwargs
        )
    X, W = MOMENTS_QUADS[distribution, order]

    if k_loc in distribution._mom_cache:
        return distribution._mom_cache[k_loc]

    out = float(n.sum(n.prod(X.T**k_loc, 1) * W))
    distribution._mom_cache[k_loc] = out
    return out
//  [python]
def approximate_moment(
    distribution, k_loc, order=100000, rule="clenshaw_curtis", **kwargs
):
    """
    Approximation method for estimation of raw statistical moments.

    Uses quadrature integration to estimate the values.

    Args:
        distribution (Distribution):
            Distribution domain with dim=len(distribution)
        k_loc (Sequence[int, ...]):
            The exponents of the moments of interest with ``shape == (dim,)``.
        order (int):
            The quadrature order used in approximation.
        rule (str):
            Quadrature rule for integrating moments.
        kwargs:
            Extra args passed to :func:`chaospy.generate_quadrature`.

    Examples:
        >>> distribution = chaospy.Uniform(1, 4)
        >>> round(chaospy.approximate_moment(distribution, (1,)), 4)
        2.5
        >>> round(chaospy.approximate_moment(distribution, (2,)), 4)
        7.0

    """
    assert isinstance(distribution, chaospy.Distribution)
    k_loc = numpy.asarray(k_loc)
    if len(distribution) > 1:
        assert (
            not distribution.stochastic_dependent
        ), "Dependent distributions does not support moment approximation."
        assert len(k_loc) == len(distribution), "incorrect size of exponents"
        return numpy.prod(
            [
                approximate_moment(
                    distribution[idx], (k_loc[idx],), order=order, rule=rule, **kwargs
                )
                for idx in range(len(distribution))
            ],
            axis=0,
        )

    k_loc = tuple(k_loc.tolist())
    assert all(
        [isinstance(k, int) for k in k_loc]
    ), "exponents must be integers: %s found" % type(k_loc[0])

    order = int(1e5 if order is None else order)
    if (distribution, order) not in MOMENTS_QUADS:
        MOMENTS_QUADS[distribution, order] = chaospy.generate_quadrature(
            order, distribution, rule=rule, **kwargs
        )
    X, W = MOMENTS_QUADS[distribution, order]

    if k_loc in distribution._mom_cache:
        return distribution._mom_cache[k_loc]

    out = float(numpy.sum(numpy.prod(X.T**k_loc, 1) * W))
    distribution._mom_cache[k_loc] = out
    return out
//  [python]
def generate_pyramid_anchors(logger, cf):
    """Generate anchors at different levels of a feature pyramid. Each scale
    is associated with a level of the pyramid, but each ratio is used in
    all levels of the pyramid.

    from configs:
    :param scales: cf.RPN_ANCHOR_SCALES , e.g. [4, 8, 16, 32]
    :param ratios: cf.RPN_ANCHOR_RATIOS , e.g. [0.5, 1, 2]
    :param feature_shapes: cf.BACKBONE_SHAPES , e.g.  [array of shapes per feature map] [80, 40, 20, 10, 5]
    :param feature_strides: cf.BACKBONE_STRIDES , e.g. [2, 4, 8, 16, 32, 64]
    :param anchors_stride: cf.RPN_ANCHOR_STRIDE , e.g. 1
    :return anchors: (N, (y1, x1, y2, x2, (z1), (z2)). All generated anchors in one array. Sorted
    with the same order of the given scales. So, anchors of scale[0] come first, then anchors of scale[1], and so on.
    """
    scales = cf.rpn_anchor_scales
    ratios = cf.rpn_anchor_ratios
    feature_shapes = cf.backbone_shapes
    anchor_stride = cf.rpn_anchor_stride
    pyramid_levels = cf.pyramid_levels
    feature_strides = cf.backbone_strides

    anchors = []
    logger.info("feature map shapes: {}".format(feature_shapes))
    logger.info("anchor scales: {}".format(scales))

    expected_anchors = [np.prod(feature_shapes[ii]) * len(ratios) * len(scales['xy'][ii]) for ii in pyramid_levels]

    for lix, level in enumerate(pyramid_levels):
        if len(feature_shapes[level]) == 2:
            anchors.append(generate_anchors(scales['xy'][level], ratios, feature_shapes[level],
                                            feature_strides['xy'][level], anchor_stride))
        else:
            anchors.append(generate_anchors_3D(scales['xy'][level], scales['z'][level], ratios, feature_shapes[level],
                                            feature_strides['xy'][level], feature_strides['z'][level], anchor_stride))

        logger.info("level {}: built anchors {} / expected anchors {} ||| total build {} / total expected {}".format(
            level, anchors[-1].shape, expected_anchors[lix], np.concatenate(anchors).shape, np.sum(expected_anchors)))

    out_anchors = np.concatenate(anchors, axis=0)
    return out_anchors
//  [python]
def forward(self, x):
        """
        Args:
            x (CTensor): the input tensor
        Returns:
            the result CTensor
        """
        self.shape = list(x.shape())
        shape, axis = self.shape, self.axis
        # the axis must be within this range (0, r-1)
        assert axis <= len(
            shape) - 1 or axis >= 0, "the axis must be within (0, %d-1)" % len(
                shape)
        # calculate the new shape
        new_shape = (1, int(n.prod(shape))) if axis == 0 else (
            int(n.prod(shape[0:axis]).astype(int)),
            int(n.prod(shape[axis:]).astype(int)))
        y = singa.Reshape(x, new_shape)
        return y
//  [python]
def forward(self, x):
        """
        Args:
            x (CTensor): the input tensor
        Returns:
            the result CTensor
        """
        self.shape = list(x.shape())
        shape, axis = self.shape, self.axis
        # the axis must be within this range (0, r-1)
        assert axis <= len(
            shape) - 1 or axis >= 0, "the axis must be within (0, %d-1)" % len(
                shape)
        # calculate the new shape
        new_shape = (1, int(np.prod(shape))) if axis == 0 else (
            int(np.prod(shape[0:axis]).astype(int)),
            int(np.prod(shape[axis:]).astype(int)))
        y = singa.Reshape(x, new_shape)
        return y
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 1
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 1
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = numpy.eye(nbasis, dtype=numpy.complex128)
    # print(I.shape)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    # print(I.shape)
    # exit()

    # Select lowest energy states for trial
    trial = I[:, :nalpha].conj()
    trial.imag[:] = 0
    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = n.prod(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    print(grid.shape)
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for iq, q in enumerate(qgrid):
        Gtrace_direct = 0
        # for a in range(nalpha):
        # compute \sum_G f_G(q)
        for i, k in enumerate(grid):
            kpq = k + q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                Gtrace_direct += trial[ikpq, 0] * Gh[0, i]
        fq_direct[iq] = Gtrace_direct

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    trial_pq = trial[:, 0].copy()
    Gh_pq = Gh[0, :].copy()

    fq_conv += nqgrid * convolve(Gh_pq, numpy.flip(trial_pq), mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_pq, numpy.flip(trial_pq)).ravel())

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 1
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 1
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = numpy.eye(nbasis, dtype=numpy.complex128)
    # print(I.shape)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    # print(I.shape)
    # exit()

    # Select lowest energy states for trial
    trial = I[:, :nalpha].conj()
    trial.imag[:] = 0
    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = numpy.prod(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    print(grid.shape)
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for iq, q in enumerate(qgrid):
        Gtrace_direct = 0
        # for a in range(nalpha):
        # compute \sum_G f_G(q)
        for i, k in enumerate(grid):
            kpq = k + q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                Gtrace_direct += trial[ikpq, 0] * Gh[0, i]
        fq_direct[iq] = Gtrace_direct

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    trial_pq = trial[:, 0].copy()
    Gh_pq = Gh[0, :].copy()

    fq_conv += nqgrid * convolve(Gh_pq, numpy.flip(trial_pq), mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_pq, numpy.flip(trial_pq)).ravel())

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def backward(self, dy):
        """
        backward of ReduceSum
        Args:
            dy (CTensor): gradient tensor.
        Returns:
            the gradient tensor over input tensor.
        """
        x_shape, x = self.cache
        dy = singa.Reshape(dy, x_shape)
        scale = n.prod(x_shape) / n.prod(x.shape())
        mask = singa.Tensor(list(x.shape()), x.device())
        mask.SetFloatValue(scale)
        dy = singa.__mul__(mask, dy)
        return dy
//  [python]
def backward(self, dy):
        """
        backward of ReduceSum
        Args:
            dy (CTensor): gradient tensor.
        Returns:
            the gradient tensor over input tensor.
        """
        x_shape, x = self.cache
        dy = singa.Reshape(dy, x_shape)
        scale = np.prod(x_shape) / np.prod(x.shape())
        mask = singa.Tensor(list(x.shape()), x.device())
        mask.SetFloatValue(scale)
        dy = singa.__mul__(mask, dy)
        return dy
//  [python]
def size(self):
        """Numpy-style attribute giving the total dataset size"""
        return n.prod(self.shape, dtype=n.intp)
//  [python]
def size(self):
        """Numpy-style attribute giving the total dataset size"""
        return numpy.prod(self.shape, dtype=numpy.intp)
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 1
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 1
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = numpy.eye(nbasis, dtype=numpy.complex128)
    # print(I.shape)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    # print(I.shape)
    # exit()

    # Select lowest energy states for trial
    trial = I[:, :nalpha].conj()
    trial.imag[:] = 0

    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = numpy.multiply.reduce(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    print(grid.shape)
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for iq, q in enumerate(qgrid):
        Gtrace_direct = 0
        # for a in range(nalpha):
        # compute \sum_G f_G(q)
        for i, k in enumerate(grid):
            # kmq = q-k
            kmq = k - q
            ikmq = lookup(kmq, grid)
            # idx = numpy.argwhere(basis == q-k)
            # print(idx, ikmq)
            if ikmq is not None:
                Gtrace_direct += trial[i, 0] * Gh[0, ikmq]
        fq_direct[iq] = Gtrace_direct
    print(f"trial[igmq,0] = {trial[:, (0)]}")
    print(f"Gh[0,i] = {Gh[(0), :]}")
    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    trial_pq = trial[:, 0].copy()
    Gh_pq = Gh[0, :].copy()

    # \sum_G f(G-Q) g(G)
    # -G-Q -> G'
    # \sum_G g(-G-Q) f(-G) = \sum_G g(G) f(G+Q)
    # trial_pq = numpy.conj(trial_pq)
    fq_conv += nqgrid * convolve(Gh_pq, numpy.flip(trial_pq), mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_pq, numpy.flip(trial_pq)).ravel())

    # for i in range(qmesh[0]):
    #     target = fq_conv[i]
    #     print(numpy.abs(target))
    #     if (numpy.abs(target) > 1e-8):
    #         # print(target)
    #         idx = numpy.argwhere(numpy.abs(fq_direct - target) < 1e-8)
    #         # print(i, idx)

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 1
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 1
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = numpy.eye(nbasis, dtype=numpy.complex128)
    # print(I.shape)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    # print(I.shape)
    # exit()

    # Select lowest energy states for trial
    trial = I[:, :nalpha].conj()
    trial.imag[:] = 0

    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = numpy.prod(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    print(grid.shape)
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for iq, q in enumerate(qgrid):
        Gtrace_direct = 0
        # for a in range(nalpha):
        # compute \sum_G f_G(q)
        for i, k in enumerate(grid):
            # kmq = q-k
            kmq = k - q
            ikmq = lookup(kmq, grid)
            # idx = numpy.argwhere(basis == q-k)
            # print(idx, ikmq)
            if ikmq is not None:
                Gtrace_direct += trial[i, 0] * Gh[0, ikmq]
        fq_direct[iq] = Gtrace_direct
    print(f"trial[igmq,0] = {trial[:, (0)]}")
    print(f"Gh[0,i] = {Gh[(0), :]}")
    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    trial_pq = trial[:, 0].copy()
    Gh_pq = Gh[0, :].copy()

    # \sum_G f(G-Q) g(G)
    # -G-Q -> G'
    # \sum_G g(-G-Q) f(-G) = \sum_G g(G) f(G+Q)
    # trial_pq = numpy.conj(trial_pq)
    fq_conv += nqgrid * convolve(Gh_pq, numpy.flip(trial_pq), mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_pq, numpy.flip(trial_pq)).ravel())

    # for i in range(qmesh[0]):
    #     target = fq_conv[i]
    #     print(numpy.abs(target))
    #     if (numpy.abs(target) > 1e-8):
    #         # print(target)
    #         idx = numpy.argwhere(numpy.abs(fq_direct - target) < 1e-8)
    #         # print(i, idx)

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def forward(self, x):
        """
        forward of ReduceMean
        Args:
            x (CTensor): input tensor.
        Returns:
            the output CTensor.
        """
        _x = tensor.from_raw_tensor(x)
        x_shape = list(_x.shape)
        # handle the special axes
        if self.axes is None:
            self.axes = [i for i in range(len(x_shape))]  # axes = None
        else:
            self.axes = [i if i >= 0 else len(x_shape) + i for i in self.axes
                        ]  # axes has negative
        self.axes.sort(reverse=True)
        for axis in self.axes:
            _x = tensor.sum(_x, axis)
            x_shape[axis] = 1
        if self.keepdims == 1:
            _x = tensor.reshape(_x, x_shape)
        self.cache = (x_shape, x)
        scale = n.prod(x_shape) / n.prod(x.shape())
        self.scale = scale
        _x = singa.MultFloat(_x.data, scale)
        return _x
//  [python]
def forward(self, x):
        """
        forward of ReduceMean
        Args:
            x (CTensor): input tensor.
        Returns:
            the output CTensor.
        """
        _x = tensor.from_raw_tensor(x)
        x_shape = list(_x.shape)
        # handle the special axes
        if self.axes is None:
            self.axes = [i for i in range(len(x_shape))]  # axes = None
        else:
            self.axes = [i if i >= 0 else len(x_shape) + i for i in self.axes
                        ]  # axes has negative
        self.axes.sort(reverse=True)
        for axis in self.axes:
            _x = tensor.sum(_x, axis)
            x_shape[axis] = 1
        if self.keepdims == 1:
            _x = tensor.reshape(_x, x_shape)
        self.cache = (x_shape, x)
        scale = np.prod(x_shape) / np.prod(x.shape())
        self.scale = scale
        _x = singa.MultFloat(_x.data, scale)
        return _x
//  [python]
def _init_bilinear(self, _, arr):
        weight = np.zeros(n.prod(arr.shape), dtype='float32')
        shape = arr.shape
        f = np.ceil(shape[3] / 2.)
        c = (2 * f - 1 - f % 2) / (2. * f)
        for i in range(n.prod(shape)):
            x = i % shape[3]
            y = (i / shape[3]) % shape[2]
            weight[i] = (1 - abs(x / f - c)) * (1 - abs(y / f - c))
        arr[:] = weight.reshape(shape)
//  [python]
def _init_bilinear(self, _, arr):
        weight = np.zeros(np.prod(arr.shape), dtype='float32')
        shape = arr.shape
        f = np.ceil(shape[3] / 2.)
        c = (2 * f - 1 - f % 2) / (2. * f)
        for i in range(np.prod(shape)):
            x = i % shape[3]
            y = (i / shape[3]) % shape[2]
            weight[i] = (1 - abs(x / f - c)) * (1 - abs(y / f - c))
        arr[:] = weight.reshape(shape)
//  [python]
def forward(self, x):
        """
        forward propogation of GlobalAveragePool
        Args:
            x (CTensor): the input tensor
        Returns:
            CTensor, the output
        """
        if training:
            self.mask = singa.Tensor(x.shape(), x.device())

        shape = list(x.shape())

        # (N x C x H x W) for channels_first
        if self.data_format == 'channels_first':
            axes = tuple(i for i in range(2, len(shape)))
            self.shape_divisor = 1 / n.prod(shape[2:])
        else:  # (N x H x W x C) for channels_last
            axes = tuple(i for i in range(1, len(shape) - 1))
            self.shape_divisor = 1 / n.prod(shape[1:-1])

        # output shape
        # (N x C x 1 x 1) for channels_first
        # (N x 1 x 1 x C) for channels_last
        for i in axes:
            shape[i] = 1

        x = tensor.from_raw_tensor(x)
        x = tensor.sum(x, axis=axes)
        x = tensor.reshape(x, shape)
        return singa.MultFloat(x.data, self.shape_divisor)
//  [python]
def forward(self, x):
        """
        forward propogation of GlobalAveragePool
        Args:
            x (CTensor): the input tensor
        Returns:
            CTensor, the output
        """
        if training:
            self.mask = singa.Tensor(x.shape(), x.device())

        shape = list(x.shape())

        # (N x C x H x W) for channels_first
        if self.data_format == 'channels_first':
            axes = tuple(i for i in range(2, len(shape)))
            self.shape_divisor = 1 / np.prod(shape[2:])
        else:  # (N x H x W x C) for channels_last
            axes = tuple(i for i in range(1, len(shape) - 1))
            self.shape_divisor = 1 / np.prod(shape[1:-1])

        # output shape
        # (N x C x 1 x 1) for channels_first
        # (N x 1 x 1 x C) for channels_last
        for i in axes:
            shape[i] = 1

        x = tensor.from_raw_tensor(x)
        x = tensor.sum(x, axis=axes)
        x = tensor.reshape(x, shape)
        return singa.MultFloat(x.data, self.shape_divisor)
//  [python]
def _init_weight(self, _, arr):
        nout = arr.shape[0]
        nin = np.prod(arr.shape[1:])
        if self.rand_type == "uniform":
            tmp = random.uniform(-1.0, 1.0, shape=(nout, nin)).asnumpy()
        elif self.rand_type == "normal":
            tmp = random.normal(0.0, 1.0, shape=(nout, nin)).asnumpy()
        u, _, v = np.linalg.svd(tmp, full_matrices=False) # pylint: disable=invalid-name
        if u.shape == tmp.shape:
            res = u
        else:
            res = v
        res = self.scale * res.reshape(arr.shape)
        arr[:] = res
//  [python]
def _init_weight(self, _, arr):
        weight = np.zeros(np.prod(arr.shape), dtype='float32')
        shape = arr.shape
        f = np.ceil(shape[3] / 2.)
        c = (2 * f - 1 - f % 2) / (2. * f)
        for i in range(np.prod(shape)):
            x = i % shape[3]
            y = (i / shape[3]) % shape[2]
            weight[i] = (1 - abs(x / f - c)) * (1 - abs(y / f - c))
        arr[:] = weight.reshape(shape)
//  [python]
def prepare_collapsed_input_and_dropout(inputs, dropout):
  assert len(inputs) == 1, "Multiple inputs not yet implemented"
  inp = inputs[0]
  shape = inp.get_shape()
  if len(shape) == 4:
    dim = int(n.prod(shape[1:4]))
    inp = tf.reshape(inp, [-1, dim])
  else:
    dim = int(shape[-1])
  if dropout != 0.0:
    keep_prob = 1.0 - dropout
    inp = tf.nn.dropout(inp, keep_prob)
  return inp, dim
//  [python]
def prepare_collapsed_input_and_dropout(inputs, dropout):
  assert len(inputs) == 1, "Multiple inputs not yet implemented"
  inp = inputs[0]
  shape = inp.get_shape()
  if len(shape) == 4:
    dim = int(numpy.prod(shape[1:4]))
    inp = tf.reshape(inp, [-1, dim])
  else:
    dim = int(shape[-1])
  if dropout != 0.0:
    keep_prob = 1.0 - dropout
    inp = tf.nn.dropout(inp, keep_prob)
  return inp, dim
//  [python]
def _init_weight(self, name, arr):
        shape = arr.shape
        hw_scale = 1.
        if len(shape) < 2:
            raise ValueError('Xavier initializer cannot be applied to vector {0}. It requires at'
                             ' least 2D.'.format(name))
        if len(shape) > 2:
            hw_scale = np.prod(shape[2:])
        fan_in, fan_out = shape[1] * hw_scale, shape[0] * hw_scale
        factor = 1.
        if self.factor_type == "avg":
            factor = (fan_in + fan_out) / 2.0
        elif self.factor_type == "in":
            factor = fan_in
        elif self.factor_type == "out":
            factor = fan_out
        else:
            raise ValueError("Incorrect factor type")
        scale = np.sqrt(self.magnitude / factor)
        if self.rnd_type == "uniform":
            random.uniform(-scale, scale, out=arr)
        elif self.rnd_type == "gaussian":
            random.normal(0, scale, out=arr)
        else:
            raise ValueError("Unknown random type")
//  [python]
def indxmap_diff(Nd):
    """
    Preindixing for rapid image gradient.
    Diff(x) = x.flat[d_indx[0]] - x.flat
    Diff_t(x) =  x.flat[dt_indx[0]] - x.flat

    :param Nd: the dimension of the image
    :type Nd: tuple with integers
    :return: d_indx: image gradient
    :return:  dt_indx:  the transpose of the image gradient
    :rtype: d_indx: lists with numpy ndarray
    :rtype: dt_indx: lists with numpy ndarray
    """

    ndims = len(Nd)
    Ndprod = reduce(lambda x, y: x * y, Nd)
    mylist = numpy.arange(0, Ndprod).astype(numpy.int32)
    mylist = numpy.reshape(mylist, Nd)
    d_indx = []
    dt_indx = []
    for pp in range(0, ndims):
        d_indx = d_indx + [ numpy.reshape(   numpy.roll(  mylist, +1 , pp  ), (Ndprod,)  ,order='C').astype(numpy.int32) ,]
        dt_indx = dt_indx + [ numpy.reshape(   numpy.roll(  mylist, -1 , pp  ) , (Ndprod,) ,order='C').astype(numpy.int32) ,]

    return d_indx,  dt_indx
//  [python]
def indxmap_diff(Nd):
    """
    Preindixing for rapid image gradient.
    Diff(x) = x.flat[d_indx[0]] - x.flat
    Diff_t(x) =  x.flat[dt_indx[0]] - x.flat

    :param Nd: the dimension of the image
    :type Nd: tuple with integers
    :return: d_indx: image gradient
    :return:  dt_indx:  the transpose of the image gradient
    :rtype: d_indx: lists with numpy ndarray
    :rtype: dt_indx: lists with numpy ndarray
    """

    ndims = len(Nd)
    Ndprod = numpy.prod(Nd)
    mylist = numpy.arange(0, Ndprod).astype(numpy.int32)
    mylist = numpy.reshape(mylist, Nd)
    d_indx = []
    dt_indx = []
    for pp in range(0, ndims):
        d_indx = d_indx + [ numpy.reshape(   numpy.roll(  mylist, +1 , pp  ), (Ndprod,)  ,order='C').astype(numpy.int32) ,]
        dt_indx = dt_indx + [ numpy.reshape(   numpy.roll(  mylist, -1 , pp  ) , (Ndprod,) ,order='C').astype(numpy.int32) ,]

    return d_indx,  dt_indx
//  [python]
def create_csr(uu, kk, Kd, Jd, M):
    Jprod = n.prod(Jd)
    csrdata =uu.ravel(order='C')

    Jdprod = n.prod(Jd)
    rowptr = n.arange(0, (M+1)*Jdprod, Jdprod)
    colindx =kk.ravel(order='C')

    csrshape = (M, n.prod(Kd))

    csr = scipy.sparse.csr_matrix((csrdata, colindx, rowptr),
                                       shape=csrshape)
    return csr
//  [python]
def create_csr(uu, kk, Kd, Jd, M):
#     Jprod = numpy.prod(Jd)
#     mm = numpy.arange(0, M).reshape( (1, M), order='C')  # indices from 0 to M-1
#     mm_rep = numpy.ones(( Jprod, 1))
#     mm = mm * mm_rep
    #     Jprod = numpy.prod(Jd)
    csrdata =uu.ravel(order='C')#numpy.reshape(uu.T, (Jprod * M, ), order='C')

    # row indices, from 1 to M convert array to list
#     rowindx = mm.ravel(order='C') #numpy.reshape(mm.T, (Jprod * M, ), order='C')

    Jdprod = numpy.prod(Jd)
    rowptr = numpy.arange(0, (M+1)*Jdprod, Jdprod)
    # colume indices, from 1 to prod(Kd), convert array to list
    colindx =kk.ravel(order='C')#numpy.reshape(kk.T, (Jprod * M, ), order='C')

    # The shape of sparse matrix
    csrshape = (M, numpy.prod(Kd))

    # Build sparse matrix (interpolator)
#     csr = scipy.sparse.csr_matrix((csrdata, (rowindx, colindx)),
#                                        shape=csrshape)
    csr = scipy.sparse.csr_matrix((csrdata, colindx, rowptr),
                                       shape=csrshape)
#     csr.has_sorted_indices = False
#     csr.sort_indices() # sort the indices in-place
    return csr
//  [python]
def convert_batch_to_space_nd(self, op):
        """batch_to_space_nd implementation."""

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        in_expr = self.get_expr(input_tensor_idx)

        input_shape = list(input_tensor.tensor.ShapeAsNumpy())
        batch = input_shape[0]

        block_shape = list(self.get_tensor_value(input_tensors[1]))
        M = len(block_shape)

        crops = list(self.get_tensor_value(input_tensors[2]))

        # From https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/batch-to-space-n-d:
        # Reshape input to reshaped of shape
        shape1 = block_shape + [batch // reduce(lambda x, y: x * y, block_shape)] + input_shape[1:]
        reshaped = _op.reshape(in_expr, newshape=shape1)

        # Permute dimensions of reshaped to produce permuted of shape
        axes = (
            [M]
            + [axis for i in range(M) for axis in [M + i + 1, i]]
            + list(range(2 * M + 1, len(shape1)))
        )
        permuted = _op.transpose(reshaped, axes=axes)

        # Reshape permuted to produce reshaped_permuted of shape
        shape2 = [0] + [-3] * M + [-2]
        reshaped_permuted = _op.reshape(permuted, newshape=shape2)

        # Crop the start and end of dimensions [1, ..., M] of reshaped_permuted according to crops
        # to produce the output of shape:
        reshaped_permuted_shape = _infer_shape(reshaped_permuted)
        cropped = reshaped_permuted
        for axis in range(1, M + 1):
            crop = crops[axis - 1]
            if (crop != [0, 0]).all():
                indices = _op.arange(
                    _expr.const(crop[0]),
                    _expr.const(reshaped_permuted_shape[axis] - crop[1]),
                    dtype="int32",
                )
                cropped = _op.take(cropped, indices=indices, axis=axis)

        return cropped