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
karllessard merged 2 commits into from
Mar 26, 2024
Merged

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

Changes from 1 commit
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
Next Next commit
Adding casts to the if test so it passes on GPU.
  • Loading branch information
@Craigacp
Craigacp committed Mar 8, 2024
commit f46cdb6c412a73f45f6324f35d438d9d2a2e7e85
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(a1, TFloat32.class)), TInt32.class);
return Signature.builder().input("a", a1).input("b", b1).output("y", y).build();
});

Expand Down