1"""Generate a mock model for LLVM tests.
3The generated model is not a neural net - it is just a tf.function with the
4correct input and output parameters. By construction, the mock model will always
12import tensorflow
as tf
14POLICY_DECISION_LABEL =
'inlining_decision'
15POLICY_OUTPUT_SPEC =
"""
18 "logging_name":
"inlining_decision",
20 "name":
"StatefulPartitionedCall",
32# pylint: disable=g-complex-comprehension
33def get_input_signature():
34 """Returns the list of features for LLVM inlining."""
37 tf.TensorSpec(dtype=tf.int64, shape=(), name=key)
for key
in [
38 'caller_basic_block_count',
39 'caller_conditionally_executed_blocks',
41 'callee_basic_block_count',
42 'callee_conditionally_executed_blocks',
54 'call_argument_setup',
55 'load_relative_intrinsic',
56 'lowered_call_arg_setup',
57 'indirect_call_penalty',
59 'case_cluster_penalty',
61 'unsimplified_common_instructions',
64 'simplified_instructions',
66 'constant_offset_ptr_args',
69 'last_call_to_static_bonus',
72 'nested_inline_cost_estimate',
79 tf.TensorSpec(dtype=tf.float32, shape=(), name=key)
80 for key
in [
'discount',
'reward']
85 tf.TensorSpec(dtype=tf.int32, shape=(), name=key)
86 for key
in [
'step_type']
92 return POLICY_DECISION_LABEL
96 return POLICY_OUTPUT_SPEC
99 return os.path.join(path,
'output_spec.json')
103 """Build and save the mock model with the given signature"""
106 return {signature[
'output']: tf.constant(value=1, dtype=tf.int64)}
108 module.action = tf.function()(action)
109 action = {
'action': module.action.get_concrete_function(signature[
'inputs'])}
110 tf.saved_model.save(module, path, signatures=action)
113 with open(output_spec_path,
'w')
as f:
114 print(f
'Writing output spec to {output_spec_path}.')
115 f.write(signature[
'output_spec'])
127 assert len(argv) == 2
130 print(f
'Output model to: [{argv[1]}]')
135if __name__ ==
'__main__':
static void print(raw_ostream &Out, object::Archive::Kind Kind, T Val)
def get_output_spec_path(path)
def get_input_signature()
def get_output_signature()
def build_mock_model(path, signature)