From 16acfbcf22c66a220176c1b95d30e7026c7231b2 Mon Sep 17 00:00:00 2001
From: Raffi Khatchadourian <khatchad@hunter.cuny.edu>
Date: Mon, 4 Mar 2024 17:35:32 -0500
Subject: [PATCH] Add summary for `tf.keras.Model` (#158)

Add attributes that are lists of tensors.
---
 .../python/ml/test/TestTensorflow2Model.java  | 36 ++++++++++++++++
 .../data/tensorflow.xml                       | 41 +++++++++++++++++++
 .../data/tf2_test_model_attributes.py         | 15 +++++++
 .../data/tf2_test_model_attributes2.py        | 16 ++++++++
 .../data/tf2_test_model_attributes3.py        | 17 ++++++++
 .../data/tf2_test_model_attributes4.py        | 16 ++++++++
 .../data/tf2_test_model_attributes5.py        | 17 ++++++++
 .../data/tf2_test_model_attributes6.py        | 16 ++++++++
 8 files changed, 174 insertions(+)
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_model_attributes.py
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_model_attributes2.py
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_model_attributes3.py
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_model_attributes4.py
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_model_attributes5.py
 create mode 100644 com.ibm.wala.cast.python.test/data/tf2_test_model_attributes6.py

diff --git a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
index f87dbdbcb..206f9ec0d 100644
--- a/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
+++ b/com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java
@@ -990,6 +990,42 @@ public void testModelCall4()
     test("tf2_test_model_call4.py", "SequentialModel.__call__", 1, 1, 3);
   }
 
+  @Test
+  public void testModelAttributes()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_model_attributes.py", "f", 1, 1, 2);
+  }
+
+  @Test
+  public void testModelAttributes2()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_model_attributes2.py", "f", 1, 1, 2);
+  }
+
+  @Test
+  public void testModelAttributes3()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_model_attributes3.py", "f", 1, 1, 2);
+  }
+
+  @Test
+  public void testModelAttributes4()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_model_attributes4.py", "f", 1, 1, 2);
+  }
+
+  @Test
+  public void testModelAttributes5()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_model_attributes5.py", "f", 1, 1, 2);
+  }
+
+  @Test
+  public void testModelAttributes6()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test("tf2_test_model_attributes6.py", "f", 1, 1, 2);
+  }
+
   @Test
   public void testCallbacks()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
diff --git a/com.ibm.wala.cast.python.ml/data/tensorflow.xml b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
index 20e381d11..ddcda65b4 100644
--- a/com.ibm.wala.cast.python.ml/data/tensorflow.xml
+++ b/com.ibm.wala.cast.python.ml/data/tensorflow.xml
@@ -45,6 +45,8 @@
         <new def="layers" class="Lobject" />
         <putfield class="LRoot" field="layers" fieldType="LRoot" ref="x" value="layers" />
         <putfield class="LRoot" field="layers" fieldType="LRoot" ref="keras" value="layers" />
+        <new def="models" class="Lobject" />
+        <putfield class="LRoot" field="models" fieldType="LRoot" ref="keras" value="models" />
         <new def="app" class="Lobject" />
         <putfield class="LRoot" field="app" fieldType="LRoot" ref="x" value="app" />
         <new def="run" class="Ltensorflow/app/run" />
@@ -146,6 +148,9 @@
         <putfield class="LRoot" field="Input" fieldType="LRoot" ref="layers" value="Input" />
         <new def="Dense" class="Ltensorflow/keras/layers/Dense" />
         <putfield class="LRoot" field="Dense" fieldType="LRoot" ref="layers" value="Dense" />
+        <new def="Model" class="Ltensorflow/keras/models/Model" />
+        <putfield class="LRoot" field="Model" fieldType="LRoot" ref="keras" value="Model" />
+        <putfield class="LRoot" field="Model" fieldType="LRoot" ref="models" value="Model" />
         <new def="Variable" class="Ltensorflow/functions/Variable" />
         <putfield class="LRoot" field="Variable" fieldType="LRoot" ref="x" value="Variable" />
         <putfield class="LRoot" field="Variable" fieldType="LRoot" ref="variables" value="Variable" />
@@ -729,6 +734,42 @@
         </method>
       </class>
     </package>
+    <package name="tensorflow/keras/models">
+      <class name="Model" allocatable="true">
+        <method name="read_data" descriptor="()LRoot;">
+          <new def="x" class="Ltensorflow/keras/Model/attribute" />
+          <return value="x" />
+        </method>
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/keras/models/Model -->
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self inputs outputs name">
+          <new def="__call__" class="Ltensorflow/keras/models/__call__" />
+          <putfield class="LRoot" field="__call__" fieldType="LRoot" ref="arg0" value="__call__" />
+          <new def="call" class="Ltensorflow/keras/models/call" />
+          <putfield class="LRoot" field="call" fieldType="LRoot" ref="arg0" value="call" />
+          <new def="x" class="Llist" />
+          <call class="LRoot" name="read_data" descriptor="()LRoot;" type="virtual" arg0="arg0" def="xx" />
+          <putfield class="LRoot" field="0" fieldType="LRoot" ref="x" value="xx" />
+          <!-- https://www.tensorflow.org/guide/keras/transfer_learning#freezing_layers_understanding_the_trainable_attribute -->
+          <putfield class="LRoot" field="trainable_weights" fieldType="LRoot" ref="arg0" value="x" />
+          <putfield class="LRoot" field="weights" fieldType="LRoot" ref="arg0" value="x" />
+          <putfield class="LRoot" field="non_trainable_weights" fieldType="LRoot" ref="arg0" value="x" />
+          <return value="arg0" />
+        </method>
+      </class>
+      <class name="__call__" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/keras/Model#call -->
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self inputs training mask">
+          <return value="inputs" />
+        </method>
+      </class>
+      <!-- FIXME: Workaround for https://github.com/wala/ML/issues/106. -->
+      <class name="call" allocatable="true">
+        <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/keras/Model#call -->
+        <method name="do" descriptor="()LRoot;" numArgs="4" paramNames="self inputs training mask">
+          <return value="inputs" />
+        </method>
+      </class>
+    </package>
     <package name="tensorflow/keras/layers">
       <class name="Dense" allocatable="true">
         <!-- https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/keras/layers/Dense -->
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes.py b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes.py
new file mode 100644
index 000000000..19848e16a
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes.py
@@ -0,0 +1,15 @@
+# From https://github.com/tensorflow/tensorflow/issues/14359#issue-272179775
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+a = tf.keras.layers.Input(shape=(64,))
+b = tf.keras.layers.Dense(5)(a)
+model = tf.keras.models.Model(a, b)
+
+for i in model.trainable_weights:
+    f(i)
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes2.py b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes2.py
new file mode 100644
index 000000000..65348e9e2
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes2.py
@@ -0,0 +1,16 @@
+# From https://github.com/tensorflow/tensorflow/issues/14359#issue-272179775
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+inputs = tf.keras.Input(shape=(3,))
+x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs)
+outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(x)
+model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+for i in model.trainable_weights:
+    f(i)
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes3.py b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes3.py
new file mode 100644
index 000000000..86b2c84b9
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes3.py
@@ -0,0 +1,17 @@
+# From https://github.com/tensorflow/tensorflow/issues/14359#issue-272179775
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+inputs = tf.keras.Input(shape=(3,))
+x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs)
+outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(x)
+model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+# From https://www.tensorflow.org/guide/keras/transfer_learning#freezing_layers_understanding_the_trainable_attribute
+for i in model.weights:
+    f(i)
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes4.py b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes4.py
new file mode 100644
index 000000000..800b85c76
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes4.py
@@ -0,0 +1,16 @@
+# From https://github.com/tensorflow/tensorflow/issues/14359#issue-272179775
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+a = tf.keras.layers.Input(shape=(64,))
+b = tf.keras.layers.Dense(5)(a)
+model = tf.keras.models.Model(a, b)
+
+# From https://www.tensorflow.org/guide/keras/transfer_learning#freezing_layers_understanding_the_trainable_attribute.
+for i in model.weights:
+    f(i)
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes5.py b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes5.py
new file mode 100644
index 000000000..71cda261f
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes5.py
@@ -0,0 +1,17 @@
+# From https://github.com/tensorflow/tensorflow/issues/14359#issue-272179775
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+inputs = tf.keras.Input(shape=(3,))
+x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs)
+outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(x)
+model = tf.keras.Model(inputs=inputs, outputs=outputs)
+
+# From https://www.tensorflow.org/guide/keras/transfer_learning#freezing_layers_understanding_the_trainable_attribute
+for i in model.non_trainable_weights:
+    f(i)
diff --git a/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes6.py b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes6.py
new file mode 100644
index 000000000..d44ff1c56
--- /dev/null
+++ b/com.ibm.wala.cast.python.test/data/tf2_test_model_attributes6.py
@@ -0,0 +1,16 @@
+# From https://github.com/tensorflow/tensorflow/issues/14359#issue-272179775
+
+import tensorflow as tf
+
+
+def f(a):
+    pass
+
+
+a = tf.keras.layers.Input(shape=(64,))
+b = tf.keras.layers.Dense(5)(a)
+model = tf.keras.models.Model(a, b)
+
+# From https://www.tensorflow.org/guide/keras/transfer_learning#freezing_layers_understanding_the_trainable_attribute.
+for i in model.non_trainable_weights:
+    f(i)