This is an archive of the discontinued LLVM Phabricator instance.

	Show All 40 Lines
	def build_mock_model(path):			def build_mock_model(path):
	"""Build and save the mock model with the given signature."""			"""Build and save the mock model with the given signature."""
	module = tf.Module()			module = tf.Module()
	# We have to set this useless variable in order for the TF C API to correctly			# We have to set this useless variable in order for the TF C API to correctly
	# intake it			# intake it
	module.var = tf.Variable(0, dtype=tf.int64)			module.var = tf.Variable(0, dtype=tf.int64)

	def action(*inputs):			def action(*inputs):
	result = tf.math.argmax(tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var			result = tf.math.argmax(
				tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
	return {POLICY_DECISION_LABEL: result}			return {POLICY_DECISION_LABEL: result}
	module.action = tf.function()(action)			module.action = tf.function()(action)
	action = {			action = {
	'action': module.action.get_concrete_function(get_input_signature())			'action': module.action.get_concrete_function(get_input_signature())
	}			}
	tf.saved_model.save(module, path, signatures=action)			tf.saved_model.save(module, path, signatures=action)
	output_spec_path = get_output_spec_path(path)			output_spec_path = get_output_spec_path(path)
	with open(output_spec_path, 'w') as f:			with open(output_spec_path, 'w') as f:
	Show All 12 Lines

[mlgo][nfc] regalloc test model generator: prep for TFLiteClosedPublic