Topi.concatenate produces negative iteration bounds

When used in conjunction with tensors that have a dimension of size 1 topi.concatenate will result in lowered code that has a negative (invalid) iteration bound.

Example:

import tvm
from tvm import relay
from tvm import te
import topi

shape = (4, 2, 3)
shape_one_layer = (4, 3)

data = te.placeholder(shape)

l0 = te.compute(
        shape_one_layer,
        lambda x, y: data[x, 0, y] + 1,
        name='layer0')

l1 = te.compute(
        shape_one_layer,
        lambda x, y: data[x, 1, y] + 1,
        name='layer1')

l0e = topi.expand_dims(l0, 1, 1) # new shape (4, 1, 3)
l1e = topi.expand_dims(l1, 1, 1) # new shape (4, 1, 3)

result = topi.concatenate((l0e, l1e), axis=1) # shape (4, 2, 3)
s = te.create_schedule(result.op)
print(tvm.lower(s, [data]))

produces:

  allocate(T_concat, float32, [24]) {
    for (x: int32, 0, 4) {
      for (y: int32, 0, 3) {
        layer1[((x*3) + y)] = ((float32*)placeholder_2[(((x*6) + y) + 3)]) + 1f32)
      }
    }
    for (x_1: int32, 0, 4) {
      for (y_1: int32, 0, 3) {
        layer0[((x_1*3) + y_1)] = ((float32*)placeholder_2[((x_1*6) + y_1)]) + 1f32)
      }
    }
    for (ax0: int32, 0, 4) {
      for (ax1: int32, 0, 2) {
        for (ax2: int32, 0, 3) {
          T_concat[(((ax0*6) + (ax1*3)) + ax2)] = @tvm_if_then_else((1 <= ax1), (float32*)layer1[((((ax0*3) + (ax1*3)) + ax2) - 3)]), (float32*)layer0[(((ax0*3) + (ax1*3)) + ax2)]), dtype=float32, type="pure_intrin")
        }
      }
    }
  }
}

The negative bound occurs on the line where T_concat is computed, in the way layer1 is accessed: layer1[((((ax0*3) + (ax1*3)) + ax2) - 3)]. When starting iteration all three variables are 0, resulting in an index of -3.

A workaround is to use if-then-else directly:

result = te.compute(
        shape,
        lambda x, t, y: te.if_then_else(t == 0, l0e[x, 0, y], l1e[x, 0, y])) # shape (4, 2, 3)

s = te.create_schedule(result.op)
print(tvm.lower(s, [data]))

produces the correct iteration bounds:

compute[(((x_2*6) + (t*3)) + y_2)] = @tvm_if_then_else((t == 0), (float32*)layer0[((x_2*3) + y_2)]), (float32*)layer1[((x_2*3) + y_2)]), dtype=float32, type="pure_intrin")

I don’t know the source of this problem. Opinions are welcome. Also, is there any better way of concatenating tensors in TVM?