What Relay op would allow for splitting a tensor along an axis?

Questions
#1

I’m curious if someone might be able to guide me as to how I might be able to implement a Relay equivalent to the Keras TimeDistributed layer. I need to import a Keras model that features it and I am willing to extend the Keras importer to support this.

Would it be correct to say that the TVM equivalent would split a tensor on the second dimension, apply the passed layer to each of the resulting tensors, and then concatenate them back together? If so, which operator should I use to implement the splitting? I don’t have much experience with Keras and, as I said, I would be very glad to PR support for TimeDistributed to the Keras importer.

#2

@altanh pointed out to me that it should be possible to use the split operator for what I proposed. I will see if I can add support for TimeDistributed into the Keras importer