TE compute about partial sum

Questions
1

Hi,I have a question about te compute.I have tried to build a compute for partial sum with Python.However, here comes a bug that I can’t fixed.Have I miss some update for the useage of this compute?

import tvm
from tvm import te

def compress_and_merge():
    input_shape = (17, 32, 32)
    output_shape = (5, 32, 32)

    input_tensor = te.placeholder(input_shape, name='input_tensor', dtype='float32')
    output_tensor = te.placeholder(output_shape, name='output_tensor', dtype='float32')

    # define the compute
    def compute(i, j, k):
        # judge if i is the last dimension of axis 0
        is_last = i == 16

        if is_last:
            # keep the last layer remained
            return input_tensor[i, j, k]
        else:
            # conpress one time with every four layers. 
            x = i // 4
            return te.sum(input_tensor[x*4:(x+1)*4, j, k], axis=0)

    # description
    output = te.compute(output_shape, compute, name='output')

    return output.op.body[0]


with tvm.target.Target('llvm'):
    # graph
    stmt = compress_and_merge()

print(stmt)

I want to make A(17,32,32) into B(5,32,32).The first four layers is compressed from A(1,32,32) to A(16,32,32),and each layer is compressed from four layers.Forn

My complier gave me an error message. “Value Error:don’t know how to convert type<class ‘slice’> to object” Does it mean “te.sum(input_tensor[x*4:(x+1)*4, j, k], axis=0)” is illegal?