Losing tensors in datasets

[khatchad]

Description

Currently, we cannot track tensors contained in tf.data.Datasets:

import tensorflow as tf


def add(a, b):
  return a + b


dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])

for element in dataset:
    c = add(element, element)

We get:

Oct 11, 2023 11:09:48 AM com.ibm.wala.cast.python.ml.test.TestTensorflowModel testTf2
INFO: Tensor analysis: answer:

Regression

List Literals

Seems to work just fine with lists:

import tensorflow as tf


def add(a, b):
  return a + b


list = [tf.ones([1, 2]), tf.ones([2, 2])]

for element in list:
    c = add(element, element)

We get:

Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.client.PythonTensorAnalysisEngine getDataflowSources
INFO: Added dataflow source [Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@103 ], v5]:[Empty].
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.client.PythonTensorAnalysisEngine getDataflowSources
INFO: Added dataflow source [Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@97 ], v5]:[Empty].
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel testTf2
INFO: Tensor analysis: answer:
[SITE_IN_NODE{<Code body of function Lscript tf2_test_tensor_list.py>:Llist in CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]},<field 0>][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v249][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[SITE_IN_NODE{<Code body of function Lscript tf2_test_tensor_list.py>:Llist in CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]}; [Core[Root]]][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[SITE_IN_NODE{<Code body of function Lscript tf2_test_tensor_list.py>:Llist in CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]},<field 1>][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v257][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v265][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[<field global element>][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v268][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v267][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v270][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v271][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py/add> Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@117 ], v2][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list.py/add> Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@117 ], v3][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Ret-V:Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@103 ]][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Ret-V:Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@97 ]][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@103 ], v5][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list.py.do()LRoot;@97 ], v5][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]

Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.InstanceFieldKey
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.cast.ipa.callgraph.ReflectedFieldPointerKey$3
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.InstanceFieldKey
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.StaticFieldKey
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.ReturnValueKey
Oct 11, 2023 2:25:11 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.ReturnValueKey

Lists Without Literals

But, creating a list without using a literal doesn't work:

import tensorflow as tf


def add(a, b):
  return a + b


list = list()

list.append(tf.ones([1, 2]))
list.append(tf.ones([2, 2]))

for element in list:
    c = add(element, element)

We get:

Oct 11, 2023 2:51:26 PM com.ibm.wala.cast.python.ml.client.PythonTensorAnalysisEngine getDataflowSources
INFO: Added dataflow source [Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list3.py.do()LRoot;@106 ], v5]:[Empty].
Oct 11, 2023 2:51:26 PM com.ibm.wala.cast.python.ml.client.PythonTensorAnalysisEngine getDataflowSources
INFO: Added dataflow source [Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list3.py.do()LRoot;@99 ], v5]:[Empty].
Oct 11, 2023 2:51:26 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel testTf2
INFO: Tensor analysis: answer:
[Node: <Code body of function Lscript tf2_test_tensor_list3.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v264][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Ret-V:Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list3.py.do()LRoot;@106 ]][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: <Code body of function Lscript tf2_test_tensor_list3.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ], v252][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Ret-V:Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list3.py.do()LRoot;@99 ]][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list3.py.do()LRoot;@106 ], v5][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]
[Node: synthetic < PythonLoader, Ltensorflow/functions/ones, do()LRoot; > Context: CallStringContext: [ script tf2_test_tensor_list3.py.do()LRoot;@99 ], v5][{[D:Symbolic,n, D:Compound,[D:Constant,28, D:Constant,28]] of pixel}]

Oct 11, 2023 2:51:26 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.ReturnValueKey
Oct 11, 2023 2:51:26 PM com.ibm.wala.cast.python.ml.test.TestTensorflowModel lambda$0
WARNING: Encountered: class com.ibm.wala.ipa.callgraph.propagation.ReturnValueKey

[khatchad]

Potentially helpful reference: https://www.tensorflow.org/api_docs/python/tf/data/Iterator