Adding casts to the if test so it passes on GPU (#528)

tensorflow · Mar 26, 2024 · 57acdc8 · 57acdc8
1 parent ff1e816
commit 57acdc8
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/op/core/IfTest.java b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/op/core/IfTest.java
@@ -27,6 +27,7 @@
 import org.tensorflow.Session;
 import org.tensorflow.Signature;
 import org.tensorflow.op.Ops;
+import org.tensorflow.types.TFloat32;
 import org.tensorflow.types.TInt32;
 
 public class IfTest {
@@ -37,15 +38,19 @@ private static Operand<TInt32> basicIf(Ops tf, Operand<TInt32> a, Operand<TInt32
             (ops) -> {
               Operand<TInt32> a1 = ops.placeholder(TInt32.class);
               Operand<TInt32> b1 = ops.placeholder(TInt32.class);
-              return Signature.builder().input("a", a1).input("b", b1).output("y", a1).build();
+              Operand<TInt32> y = ops.identity(a1);
+              return Signature.builder().input("a", a1).input("b", b1).output("y", y).build();
             });
 
     ConcreteFunction elseBranch =
         ConcreteFunction.create(
             (ops) -> {
               Operand<TInt32> a1 = ops.placeholder(TInt32.class);
               Operand<TInt32> b1 = ops.placeholder(TInt32.class);
-              Operand<TInt32> y = ops.math.neg(b1);
+              // Casts around the math.neg operator as it's not implemented correctly for int32 in
+              // GPUs at some point between TF 2.10 and TF 2.15.
+              Operand<TInt32> y =
+                  ops.dtypes.cast(ops.math.neg(ops.dtypes.cast(b1, TFloat32.class)), TInt32.class);
               return Signature.builder().input("a", a1).input("b", b1).output("y", y).build();
             });