fix for key2value

Author: kere-nel

src/Microsoft.ML.Data/Transforms/KeyToValue.cs

@@ -505,19 +505,20 @@ public override bool SaveOnnx(OnnxContext ctx, string srcVariableName, string ds
                     // Onnx expects the input keys to be int64s. But the input data can come from an ML.NET node that
                     // may output a uint32. So cast it here to ensure that the data is treated correctly
                     opType = "Cast";
-                    var castNodeOutput = ctx.AddIntermediateVariable(NumberDataViewType.Int64, "CastNodeOutput");
+                    var srcShape = (int)ctx.RetrieveShapeOrNull(srcVariableName)[1];
+                    var castNodeOutput = ctx.AddIntermediateVariable(new VectorDataViewType(NumberDataViewType.Int64, srcShape), "CastNodeOutput");
                     var castNode = ctx.CreateNode(opType, srcVariableName, castNodeOutput, ctx.GetNodeName(opType), "");
                     var t = InternalDataKindExtensions.ToInternalDataKind(DataKind.Int64).ToType();
                     castNode.AddAttribute("to", t);
 
                     var labelEncoderOutput = dstVariableName;
                     var labelEncoderInput = srcVariableName;
                     if (TypeOutput == NumberDataViewType.Double || TypeOutput == BooleanDataViewType.Instance)
-                        labelEncoderOutput = ctx.AddIntermediateVariable(NumberDataViewType.Single, "CastNodeOutput");
+                        labelEncoderOutput = ctx.AddIntermediateVariable(new VectorDataViewType(NumberDataViewType.Single, srcShape), "CastNodeOutput");
                     else if (TypeOutput == NumberDataViewType.Int64 || TypeOutput == NumberDataViewType.UInt16 ||
                         TypeOutput == NumberDataViewType.Int32 || TypeOutput == NumberDataViewType.Int16 ||
                         TypeOutput == NumberDataViewType.UInt64 || TypeOutput == NumberDataViewType.UInt32)
-                        labelEncoderOutput = ctx.AddIntermediateVariable(TextDataViewType.Instance, "CastNodeOutput");
+                        labelEncoderOutput = ctx.AddIntermediateVariable(new VectorDataViewType(TextDataViewType.Instance, srcShape), "CastNodeOutput");
 
                     opType = "LabelEncoder";
                     var node = ctx.CreateNode(opType, castNodeOutput, labelEncoderOutput, ctx.GetNodeName(opType));

src/Microsoft.ML.Data/Transforms/ValueToKeyMappingTransformer.cs

@@ -787,7 +787,8 @@ private IEnumerable<T> GetTermsAndIds<T>(int iinfo, out long[] termIds)
             private void CastInputToString<T>(OnnxContext ctx, out OnnxNode node, out long[] termIds, string srcVariableName, int iinfo,
                 string opType, string labelEncoderOutput)
             {
-                var castOutput = ctx.AddIntermediateVariable(TextDataViewType.Instance, "castOutput");
+                var srcShape = ctx.RetrieveShapeOrNull(srcVariableName);
+                var castOutput = ctx.AddIntermediateVariable(new VectorDataViewType(TextDataViewType.Instance, (int)srcShape[1]), "castOutput");
                 var castNode = ctx.CreateNode("Cast", srcVariableName, castOutput, ctx.GetNodeName("Cast"), "");
                 var t = InternalDataKindExtensions.ToInternalDataKind(DataKind.String).ToType();
                 castNode.AddAttribute("to", t);
@@ -799,7 +800,8 @@ private void CastInputToString<T>(OnnxContext ctx, out OnnxNode node, out long[]
             private void CastInputToFloat<T>(OnnxContext ctx, out OnnxNode node, out long[] termIds, string srcVariableName, int iinfo,
                 string opType, string labelEncoderOutput)
             {
-                var castOutput = ctx.AddIntermediateVariable(NumberDataViewType.Single, "castOutput");
+                var srcShape = ctx.RetrieveShapeOrNull(srcVariableName);
+                var castOutput = ctx.AddIntermediateVariable(new VectorDataViewType(NumberDataViewType.Single, (int)srcShape[1]), "castOutput");
                 var castNode = ctx.CreateNode("Cast", srcVariableName, castOutput, ctx.GetNodeName("Cast"), "");
                 var t = InternalDataKindExtensions.ToInternalDataKind(DataKind.Single).ToType();
                 castNode.AddAttribute("to", t);
@@ -813,7 +815,7 @@ private bool SaveAsOnnxCore(OnnxContext ctx, int iinfo, ColInfo info, string src
                 long[] termIds;
                 string opType = "LabelEncoder";
                 OnnxNode castNode;
-                var labelEncoderOutput = ctx.AddIntermediateVariable(NumberDataViewType.Int64, "LabelEncoderOutput");
+                var labelEncoderOutput = ctx.AddIntermediateVariable(new VectorDataViewType(NumberDataViewType.Int64, _types[iinfo].GetValueCount()), "LabelEncoderOutput");
 
                 var type = info.TypeSrc.GetItemType();
                 if (type.Equals(TextDataViewType.Instance))