Multiple right hand sides for reduction

Hey, I am testing out various TVM pathways for doing general data processing tasks.

One thing that seemed like low hanging fruit was kmeans++. I fairly quickly hit an interesting snag that I am curious about.

I was going to attempt to do the diff-squared calculation in the reduction itself and hit an error. In order to do this I attempted to declare a commreducer that had 1 lhs argument (the summation) and 2 rhs arguments (the two elemwise items of the two vectors):

        expanded-distances-op (ast/compute
                               [n-rows n-centers]
                               (ast/tvm-fn
                                [row-idx center-idx]
                                (ast/commutative-reduce
                                 (ast/tvm-fn->commutative-reducer
                                  (ast/tvm-fn
                                   [sum row-elem center-elem]
                                   (ast/tvm-let [diff (ast-op/- row-elem center-elem)]
                                                (ast-op/+ sum (ast-op/mul diff diff))))
                                  [(float 0.0)])

                                 [{:domain [0 n-cols] :name "elem-idx"}]

                                 [(fn [elem-idx]
                                    (ast/tget dataset [row-idx elem-idx]))
                                  (fn [elem-idx]
                                    (ast/tget dataset [(ast/tget center-indexes [center-idx])
                                                       elem-idx]))]))
                               "expanded-distances")

This hit an error in expr.cc that I was very curious about:

I can in this case and in general lift a compound statement out into an intermediary but I thought the restriction that I had to have an equal number of accumulators and inputs to a reduction was very odd.

Is perhaps this first check not necessary? Are there other work arounds (such as creating a tuple) I could try beside simply creating another compute operation?

Thanks for your time, I really like the work done on the object model and am having a great time playing with TVM.