'x': tf.TensorSpec([4], tf.float32), 'y': tf.TensorSpec([5], tf.float32), 'z': tf.TensorSpec([6], tf.float32), 'a': tf.TensorSpec([1], tf.float32), 'b': tf.TensorSpec([2], tf.float32), 'c': tf.TensorSpec([3], tf.float32), }]) def f0004_dict_many_keys(self, d): return # Check a slightly more complex recursive structure. # Note that list elements can have heterogenous types. # # CHECK: func {{@[a-zA-Z_0-9]+}}( # CHECK-SAME: %arg0: tensor<1xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "x", 0]}, # CHECK-SAME: %arg1: tensor<2xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "x", 1]}, # CHECK-SAME: %arg2: tensor<3xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "y"]}) # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["f0005_more_complex_recursive_structure"] @tf.function(input_signature=[{ 'x': [tf.TensorSpec([1], tf.float32), tf.TensorSpec([2], tf.float32)], 'y': tf.TensorSpec([3], tf.float32), }]) def f0005_more_complex_recursive_structure(self, d): return if __name__ == '__main__': common.do_test(TestModule)
# # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== # RUN: %p/partially_shaped_variables | FileCheck %s # pylint: disable=missing-docstring,line-too-long from __future__ import absolute_import from __future__ import division from __future__ import print_function import tensorflow.compat.v2 as tf from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common class TestModule(tf.Module): def __init__(self): super(TestModule, self).__init__() # CHECK: "tf_saved_model.global_tensor"() {is_mutable, {{.*}} tf_saved_model.exported_names = ["v0"], type = tensor<*xf32>, value = dense<0.000000e+00> : tensor<1xf32>} : () -> () # CHECK: "tf_saved_model.global_tensor"() {is_mutable, {{.*}} tf_saved_model.exported_names = ["v1"], type = tensor<?xf32>, value = dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>} : () -> () self.v0 = tf.Variable([0.], shape=tf.TensorShape(None)) self.v1 = tf.Variable([0., 1.], shape=[None]) if __name__ == '__main__': common.do_test(TestModule, exported_names=[])
# # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== # RUN: %p/debug_info | FileCheck %s # pylint: disable=missing-docstring,line-too-long import tensorflow.compat.v2 as tf from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common class TestModule(tf.Module): @tf.function(input_signature=[ tf.TensorSpec([], tf.float32), tf.TensorSpec([], tf.float32) ]) def some_function(self, x, y): return x + y # Basic check that the debug info file is being correctly saved and loaded. # # CHECK: "tf.AddV2"{{.*}}loc(#[[LOC:.*]]) # CHECK: #[[LOC]] = loc({{.*}}callsite("{{[^"]*}}/debug_info.py{{.*}}":{{[0-9]+}}:{{[0-9]+}} if __name__ == '__main__': common.do_test(TestModule, show_debug_info=True)
import tensorflow.compat.v2 as tf from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common def mnist_model(): """Creates a MNIST model.""" model = tf.keras.models.Sequential() model.add(tf.keras.layers.Flatten()) model.add(tf.keras.layers.Dense(128, activation='relu')) model.add(tf.keras.layers.Dense(10, activation='softmax')) return model class TestModule(tf.Module): def __init__(self): super(TestModule, self).__init__() self.model = mnist_model() # CHECK: func {{@[a-zA-Z_0-9]+}}(%arg0: tensor<1x28x28x1xf32> {tf_saved_model.index_path = [0]} # CHECK: attributes {{.*}} tf_saved_model.exported_names = ["my_predict"] @tf.function(input_signature=[ tf.TensorSpec([1, 28, 28, 1], tf.float32), ]) def my_predict(self, x): return self.model(x) if __name__ == '__main__': common.do_test(TestModule, exported_names=['my_predict'])