From 8b4b8977a0db1f8dde8553898ea23303ec8eb0a2 Mon Sep 17 00:00:00 2001
From: Alex Black <blacka101@gmail.com>
Date: Wed, 3 Jul 2019 12:37:14 +1000
Subject: [PATCH] Batch Mmul v2 import fix (#41)

Signed-off-by: AlexDBlack <blacka101@gmail.com>
---
 .../nd4j/imports/graphmapper/tf/TFGraphMapper.java   |  3 ++-
 .../org/nd4j/linalg/api/ops/impl/reduce/Mmul.java    | 12 +++++++-----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/graphmapper/tf/TFGraphMapper.java b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/graphmapper/tf/TFGraphMapper.java
index c734fffc6..27f5ed693 100644
--- a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/graphmapper/tf/TFGraphMapper.java
+++ b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/imports/graphmapper/tf/TFGraphMapper.java
@@ -945,7 +945,8 @@ public class TFGraphMapper extends BaseGraphMapper<GraphDef,NodeDef,AttrValue,No
             } else if(tensorProto.getOp().equals("Assert")){
                 return org.nd4j.linalg.api.buffer.DataType.BOOL;
             }
-            log.warn("No TensorFlow descriptor found for tensor \"{}\", op \"{}\"", tensorProto.getName(), tensorProto.getOp());
+            //Not in ops.proto
+            log.debug("No TensorFlow descriptor found for tensor \"{}\", op \"{}\"", tensorProto.getName(), tensorProto.getOp());
 
             //No descriptor... try to fall back on common type attribute names
             if(!tensorProto.containsAttr("dtype") && !tensorProto.containsAttr("Tidx") && !tensorProto.containsAttr("T"))
diff --git a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/reduce/Mmul.java b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/reduce/Mmul.java
index 5abe5129c..00cad1f88 100644
--- a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/reduce/Mmul.java
+++ b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/ops/impl/reduce/Mmul.java
@@ -153,7 +153,7 @@ public class Mmul extends DynamicCustomOp {
 
     @Override
     public String[] tensorflowNames() {
-        return new String[]{"MatMul", "BatchMatMul"};
+        return new String[]{"MatMul", "BatchMatMul", "BatchMatMulV2"};
     }
 
 
@@ -171,7 +171,12 @@ public class Mmul extends DynamicCustomOp {
 
         boolean isTransposeA;
         boolean isTransposeB;
-        if(nodeDef.getOp().equalsIgnoreCase("BatchMatMul")){
+        if(nodeDef.getOp().equalsIgnoreCase("MatMul")){
+            isTransposeA = attributesForNode.get("transpose_a").getB();
+            isTransposeB = attributesForNode.get("transpose_b").getB();
+
+        } else {
+            //BatchMatMul, BatchMatMulV2
             //In practice, BatchMatMul seems to use "adj_x" and "adj_y" instead of "transpose_a" and "transpose_b"
             if(attributesForNode.containsKey("transpose_a")){
                 isTransposeA = attributesForNode.get("transpose_a").getB();
@@ -183,9 +188,6 @@ public class Mmul extends DynamicCustomOp {
             } else {
                 isTransposeB = attributesForNode.get("adj_y").getB();
             }
-        } else {
-            isTransposeA = attributesForNode.get("transpose_a").getB();
-            isTransposeB = attributesForNode.get("transpose_b").getB();
         }
         MMulTranspose mMulTranspose = MMulTranspose.builder()
                 .transposeA(isTransposeA).transposeB(isTransposeB)