Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 26, 2024
Merged

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

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f46cdb6
Adding casts to the if test so it passes on GPU.
Craigacp Mar 8, 2024
8d48f57
Fix a typo in the variable name
Craigacp Mar 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/op/core/IfTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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();
});

Expand Down