Add callback (#292)

* Upgrade ML for #280. * Adding callback test * Fixing comment * Must use at least ML 0.13.0. * Simplify test. --------- Co-authored-by: Tatiana Castro Velez <[email protected]>
ponder-lab · Dec 5, 2023 · c01ffc5 · c01ffc5
1 parent d70f0d6
commit c01ffc5
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 4 deletions.
diff --git a/edu.cuny.hunter.hybridize.core/META-INF/MANIFEST.MF b/edu.cuny.hunter.hybridize.core/META-INF/MANIFEST.MF
@@ -26,8 +26,8 @@ Import-Package: com.google.common.collect,
  com.ibm.wala.cast.python.client;version="0.1.0",
  com.ibm.wala.cast.python.ipa.callgraph;version="0.12.0",
  com.ibm.wala.cast.python.loader;version="0.1.0",
- com.ibm.wala.cast.python.ml.analysis;version="0.1.0",
- com.ibm.wala.cast.python.ml.client;version="0.1.0",
+ com.ibm.wala.cast.python.ml.analysis;version="0.13.0",
+ com.ibm.wala.cast.python.ml.client;version="0.13.0",
  com.ibm.wala.cast.python.modref;version="0.11.0",
  com.ibm.wala.cast.python.ssa,
  com.ibm.wala.cast.python.types,

diff --git a/edu.cuny.hunter.hybridize.core/tensorflow.xml b/edu.cuny.hunter.hybridize.core/tensorflow.xml
@@ -36,6 +36,9 @@
         <new def="estimator" class="Lobject" />
         <putfield class="LRoot" field="estimator" fieldType="LRoot" ref="x" value="estimator" />
 
+        <new def="distribute" class="Lobject" />
+        <putfield class="LRoot" field="distribute" fieldType="LRoot" ref="x" value="distribute" />
+
         <new def="nn" class="Lobject" />
         <putfield class="LRoot" field="nn" fieldType="LRoot" ref="x" value="nn" />
         <new def="random" class="Lobject" />
@@ -58,6 +61,9 @@
         <new def="Estimator" class="Ltensorflow/estimator/Estimator" />
         <putfield class="LRoot" field="Estimator" fieldType="LRoot" ref="estimator" value="Estimator" />
 
+        <new def="MirroredStrategy" class="Ltensorflow/distribute/MirroredStrategy" />
+        <putfield class="LRoot" field="MirroredStrategy" fieldType="LRoot" ref="distribute" value="MirroredStrategy" />
+
         <new def="inputs" class="Lobject" />
         <putfield class="LRoot" field="inputs" fieldType="LRoot" ref="estimator" value="inputs" />
 
@@ -786,6 +792,26 @@
       </class>
     </package>
 
+    <package name="tensorflow/distribute">
+      <class name="MirroredStrategy" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self devices cross_device_ops">
+          <new def="x" class="Ltensorflow/distribute/run/run" />
+          <putfield class="LRoot" field="run" fieldType="LRoot" ref="self" value="x" />
+          <return value="arg0" />
+        </method>
+      </class>
+    </package>
+
+    <package name="tensorflow/distribute/run">
+      <class name="run" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="3">
+          <getfield class="LRoot" field="0" fieldType="LRoot" ref="arg2" def="x" />
+          <call class="LRoot" name="do" descriptor="()LRoot;" type="virtual" arg0="arg1" arg1="x" numArgs="2" def="v" />
+          <return value="v" />
+        </method>
+      </class>
+    </package>
+
     <package name="tensorflow/app">
       <class name="run" allocatable="true">
         <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self main argv">
@@ -857,4 +883,4 @@
     </package>
 
   </classloader>
-</summary-spec>
+</summary-spec>
diff --git a/edu.cuny.hunter.hybridize.tests/resources/HybridizeFunction/testCallback/in/A.py b/edu.cuny.hunter.hybridize.tests/resources/HybridizeFunction/testCallback/in/A.py
@@ -0,0 +1,16 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/Strategy#example_usage_2.
+
+import tensorflow as tf
+
+strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+tensor_input = tf.constant(3.0)
+
+
+@tf.function
+def replica_fn(input):
+  return input * 2.0
+
+
+# Indirect call to replica_fun().
+result = strategy.run(replica_fn, (tensor_input,))
+print(result)
diff --git a/edu.cuny.hunter.hybridize.tests/resources/HybridizeFunction/testCallback/in/requirements.txt b/edu.cuny.hunter.hybridize.tests/resources/HybridizeFunction/testCallback/in/requirements.txt
@@ -0,0 +1 @@
+tensorflow==2.9.3
diff --git a/...ze.tests/test cases/edu/cuny/hunter/hybridize/tests/HybridizeFunctionRefactoringTest.java b/...ze.tests/test cases/edu/cuny/hunter/hybridize/tests/HybridizeFunctionRefactoringTest.java
@@ -5664,5 +5664,14 @@ public void testRecursion11() throws Exception {
 		assertTrue(f.getEntryMatchingFailure(IS_RECURSIVE).isError());
 	}
 
+	/**
+	 * Test for https://github.com/ponder-lab/Hybridize-Functions-Refactoring/issues/280.
+	 */
+	@Test
+	public void testCallback() throws Exception {
+		Function f = getFunction("replica_fn");
+		assertTrue(f.getLikelyHasTensorParameter());
+	}
+
 	// TODO: Left off at https://www.tensorflow.org/guide/function#depending_on_python_global_and_free_variables.
 }
diff --git a/hybridize.target b/hybridize.target
@@ -29,7 +29,7 @@
 				<dependency>
 					<groupId>com.ibm.wala</groupId>
 					<artifactId>com.ibm.wala.cast.python.ml</artifactId>
-					<version>0.13.0-SNAPSHOT</version>
+					<version>0.14.0-SNAPSHOT</version>
 					<type>jar</type>
 				</dependency>
 			</dependencies>