[Model] Update Mixtral to have well-formed TIR

Inside a `T.block`, loop variables may not be used, and access to them
must be done through the corresponding `T.axis.remap` output.

mlc_llm/relax_model/mixtral.py

@@ -350,14 +350,14 @@ def top2_softmax_func(
                         for j in T.unroll(2):
                             with T.block("cast"):
                                 vj = T.axis.remap("S", [j])
-                                local_top_k_f32[vj] = T.cast(local_top_k[j], "float32")
+                                local_top_k_f32[vj] = T.cast(local_top_k[vj], "float32")
                         with T.block("max"):
                             local_top_k_max[0] = T.max(local_top_k_f32[0], local_top_k_f32[1])
                         for j in T.unroll(2):
                             with T.block("output"):
                                 vj = T.axis.remap("S", [j])
                                 out[vi, vj] = T.cast(
-                                    T.exp(local_top_k_f32[j] - local_top_k_max[0])
+                                    T.exp(local_top_k_f32[vj] - local_top_k_max[0])
                                     / (
                                         T.exp(local_top_k_f32[0] - local_top_k_max[0])
                                         + T.exp(local_top_k_f32[1] - local_top_k_max[0])