Fix input shape for matmulnbits to fold input via reshape to 1d input

src/onnx/parse_matmulnbits.cpp

@@ -67,7 +67,15 @@ struct parse_matmulnbits : op_parser<parse_matmulnbits>
                            ". Actual dims: " + to_string_range(args[1]->get_shape().lens()));
 
         std::vector<size_t> expected_scales_lens{n * n_blocks_per_col};
-        if(args[2]->get_shape().lens() != expected_scales_lens)
+
+        // Reshape anything larger than 1 dimension into a 1d tensor so we can check if we have the right amount of elements. 
+        auto scale_input = args[2];
+        if(scale_input->get_shape().lens().size() > 1)
+        {
+            scale_input = info.add_instruction(make_op("reshape", {{"dims", {scale_input->get_shape().elements()}}}), scale_input);
+        }
+
+        if(scale_input->get_shape().lens() != expected_scales_lens)
             MIGRAPHX_THROW("MatMulNBits: Input scales does not match expected dims: " +
                            to_string_range(expected_scales_lens) +
                            ". Actual dims: " + to_string_range(args[2]->get_shape().lens()));