Merge pull request apache#137 from tqchen/master

Enable better TShape parsing, support long, and integer
stefanhenneking · Jun 22, 2016 · 948abff · 948abff
2 parents 9faf1af + d3a049f
commit 948abff
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/mshadow/tensor_blob.h b/mshadow/tensor_blob.h
@@ -386,7 +386,15 @@ inline std::ostream &operator<<(std::ostream &os, const TShape &shape) {
 inline std::istream &operator>>(std::istream &is, TShape &shape) {
   // get (
   while (true) {
-    char ch = is.get();
+    char ch = is.peek();
+    if (isdigit(ch)) {
+      index_t idx;
+      if (is >> idx) {
+        shape.CopyFrom(&idx, &idx + 1);
+      }
+      return is;
+    }
+    is.get();
     if (ch == '(') break;
     if (!isspace(ch)) {
       is.setstate(std::ios::failbit);
@@ -401,6 +409,9 @@ inline std::istream &operator>>(std::istream &is, TShape &shape) {
     do {
       ch = is.get();
     } while (isspace(ch));
+    if (ch == 'L') {
+      ch = is.get();
+    }
     if (ch == ',') {
       while (true) {
         ch = is.peek();