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();
  });