Fix missing callbacks

com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflowModel.java

@@ -217,6 +217,13 @@ public void testTf2()
     // once https://github.com/wala/ML/issues/106 is fixed.
     testTf2("tf2_test_model_call3.py", "SequentialModel.call", 1, 4, 3);
     testTf2("tf2_test_model_call4.py", "SequentialModel.__call__", 1, 4, 3);
+    testTf2("tf2_test_callbacks.py", "replica_fn", 1, 3, 2);
+    testTf2(
+        "tf2_test_callbacks2.py",
+        "replica_fn",
+        1,
+        3,
+        2); // NOTE: We may not wind up with 1, 3, 2, but we need the one tensor param.
   }
 
   private void testTf2(

com.ibm.wala.cast.python.ml/data/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,59 @@
       </class>
     </package>
 
+    <package name="tensorflow/distribute">
+      <class name="MirroredStrategy" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self function arguments">
+          <new def="x" class="Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function" />
+          <putfield class="LRoot" field="distribute_datasets_from_function" fieldType="LRoot" ref="self" value="x" />
+          <return value="arg0" />
+          <new def="y" class="Ltensorflow/distribute/experimental_distribute_datasets_from_function/experimental_distribute_datasets_from_function" />
+          <putfield class="LRoot" field="experimental_distribute_datasets_from_function" fieldType="LRoot" ref="self" value="y" />
+          <return value="arg1" />
+          <new def="z" class="Ltensorflow/distribute/run/run" />
+          <putfield class="LRoot" field="run" fieldType="LRoot" ref="self" value="z" />
+          <return value="arg2" />
+        </method>
+      </class>
+    </package>
+
+    <package name="tensorflow/distribute/distribute_datasets_from_function">
+      <class name="distribute_datasets_from_function" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self dataset_fn">
+          <new def="x" class="Lobject" />
+          <putfield class="LRoot" field="MirroredStrategy" fieldType="LRoot" ref="self" value="x" />
+          <putfield class="LRoot" field="$callback" fieldType="LRoot" ref="x" value="dataset_fn" />
+          <call class="LRoot" name="do" descriptor="()LRoot;" type="virtual" arg0="dataset_fn" arg1="1" numArgs="2" def="xx" />
+          <return value="xx" />
+        </method>
+      </class>
+    </package>
+
+    <package name="tensorflow/distribute/experimental_distribute_datasets_from_function">
+      <class name="experimental_distribute_datasets_from_function" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self dataset_fn">
+          <new def="x" class="Lobject" />
+          <putfield class="LRoot" field="MirroredStrategy" fieldType="LRoot" ref="self" value="x" />
+          <putfield class="LRoot" field="$callback" fieldType="LRoot" ref="x" value="dataset_fn" />
+          <call class="LRoot" name="do" descriptor="()LRoot;" type="virtual" arg0="dataset_fn" arg1="1" numArgs="2" def="xx" />
+          <return value="xx" />
+        </method>
+      </class>
+    </package>
+
+    <package name="tensorflow/distribute/run">
+      <class name="run" allocatable="true">
+        <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self fn args">
+          <new def="x" class="Lobject" />
+          <putfield class="LRoot" field="MirroredStrategy" fieldType="LRoot" ref="self" value="x" />
+          <putfield class="LRoot" field="$callback" fieldType="LRoot" ref="x" value="fn" />
+          <putfield class="LRoot" field="args" fieldType="LRoot" ref="x" value="args" />
+          <call class="LRoot" name="do" descriptor="()LRoot;" type="virtual" arg0="fn" arg1="args" numArgs="2" def="xx" />
+          <return value="xx" />
+        </method>
+      </class>
+    </package>
+
     <package name="tensorflow/app">
       <class name="run" allocatable="true">
         <method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self main argv">

com.ibm.wala.cast.python.test/data/tf2_test_callbacks.py

@@ -0,0 +1,15 @@
+# From https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/Strategy#example_usage_2.
+
+import tensorflow as tf
+
+tensor_input = tf.constant(3.0)
+
+
+@tf.function
+def replica_fn(input):
+  return input * 2.0
+
+
+# Direct call.
+result = replica_fn(tensor_input)
+print(result)

com.ibm.wala.cast.python.test/data/tf2_test_callbacks2.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)