def test_ne(self): value_1a = value.InputValue(tf.constant([1, 2]), 'a') value_1b = value.OutputValue(tf.constant([1, 2])) value_2 = value.InputValue(tf.constant([1, 3]), 'a') # We want to explicitly test __ne__. self.assertFalse(value_1a != value_1b) # pylint: disable=g-generic-assert self.assertTrue(value_1a != value_2) # pylint: disable=g-generic-assert
def operation_multipliers_from_tensor_model(
benchmark: benchmark_module.Benchmark,
tensor_model: tensor_features_model.Model, tensor_config: Dict[Text,
Any],
settings: settings_module.Settings) -> Dict[Text, float]:
"""Runs the tensor features model to get operation weight multipliers."""
# TODO(kshi): Duplicate creation of InputValue and OutputValue objects.
inputs = [
value_module.InputValue(user_input, '')
for user_input in _user_inputs(benchmark.examples[0].inputs)
]
output = value_module.OutputValue(benchmark.examples[0].output)
example_protos = collect_tensor_data.create_tf_examples(
io_example=collect_tensor_data.IOExample(
expression='',
input_values=inputs,
output_value=output,
num_inputs=benchmark.num_inputs,
operations=[]),
max_num_inputs=tensor_config['max_num_inputs'],
permute_inputs=False)
if example_protos:
example_proto = example_protos[0]
result = tensor_features_model.eval_single_example(
tensor_model, example_proto.SerializeToString())
operation_ranks = tf.argsort(
tf.argsort(tf.squeeze(result.operation_logits), stable=True))
operation_probabilities = tf.squeeze(
tf.sigmoid(result.operation_logits))
operation_names = tensor_config['operation_names']
assert len(operation_names) == len(operation_probabilities)
multipliers = {}
for index in tf.where(operation_probabilities >=
settings.tensor_model.prioritize_threshold):
chosen_op_name = operation_names[int(index)]
if settings.printing.prioritized_operations:
print('Tensor features model prioritized {}, p={}'.format(
chosen_op_name, operation_probabilities[int(index)]))
multipliers[
chosen_op_name] = settings.tensor_model.prioritize_multiplier
for index in tf.where(operation_probabilities <=
settings.tensor_model.deprioritize_threshold):
int_index = int(index)
chosen_op_name = operation_names[int_index]
if (int(operation_ranks[int_index]) >=
settings.tensor_model.max_deprioritized):
continue
if settings.printing.deprioritized_operations:
print('Tensor features model deprioritized {}, p={}, logit={}'.
format(chosen_op_name,
operation_probabilities[int_index],
result.operation_logits[0][int_index]))
multipliers[chosen_op_name] = (
settings.tensor_model.deprioritize_multiplier)
return multipliers
else:
logging.error('Failed to create example for inputs %s and output %s',
inputs, output)
return {}
def setUp(self):
super(CollectTensorDataTest, self).setUp()
self.settings = settings_module.default_settings()
operations = all_operations.get_operations()
self.unique_with_counts_operation = all_operations.find_operation_with_name(
'tf.unique_with_counts(x)', operation_list=operations)
self.indexing_operation = all_operations.find_operation_with_name(
'IndexingOperation', operation_list=operations)
self.gather_operation = all_operations.find_operation_with_name(
'tf.gather(params, indices)', operation_list=operations)
self.add_operation = all_operations.find_operation_with_name(
'tf.add(x, y)', operation_list=operations)
# Example with many operations.
in1 = value_module.InputValue([1, 1, 2, 5, 6, 5], 'in1')
in2 = value_module.InputValue([0, 10, 20, 30, 40, 50, 60, 70], 'in2')
constant_1 = value_module.ConstantValue(1)
unique = self.unique_with_counts_operation.apply([in1], self.settings)
indexed = self.indexing_operation.apply([unique, constant_1],
self.settings)
gathered = self.gather_operation.apply([in2, in1], self.settings)
self.example_value_1 = self.add_operation.apply([indexed, gathered],
self.settings)
self.assertEqual(
self.example_value_1.reconstruct_expression(),
'tf.add(tf.unique_with_counts(in1)[1], tf.gather(in2, in1))')
self.assertEqual(self.example_value_1,
value_module.OutputValue([10, 10, 21, 52, 63, 52]))
# Example with many variables and new inputs.
in3 = value_module.InputValue([1], 'in3')
in4 = value_module.InputValue([2], 'in4')
a = self.add_operation.apply([in3, new_input([10])], self.settings)
b = self.add_operation.apply([in4, in3], self.settings)
c = self.add_operation.apply([new_input([20]), in3], self.settings)
d = self.add_operation.apply([a, b], self.settings)
self.example_value_2 = self.add_operation.apply([c, d], self.settings)
self.assertEqual(
self.example_value_2.reconstruct_expression(),
'tf.add(tf.add(NEW_INPUT, in3), '
'tf.add(tf.add(in3, NEW_INPUT), tf.add(in4, in3)))')
self.assertEqual(self.example_value_2, value_module.OutputValue([35]))
def test_extract_examples_from_value(self, max_num_inputs,
expected_expressions):
actual = collect_tensor_data.extract_examples_from_value(
self.example_value_1, max_num_inputs=max_num_inputs)
# Check that all expressions are as expected.
self.assertCountEqual([example.expression for example in actual],
expected_expressions)
# Check all elements of one IOExample namedtuple. This example is at index 1
# when max_num_inputs > 1.
if max_num_inputs > 1:
expected_index_1 = collect_tensor_data.IOExample(
expression='tf.add(tf.unique_with_counts(in1)[1], in2)',
input_values=[
value_module.InputValue([1, 1, 2, 5, 6, 5], 'in1'),
value_module.InputValue([10, 10, 20, 50, 60, 50], 'in2')
],
output_value=value_module.OutputValue([10, 10, 21, 52, 63,
52]),
num_inputs=2,
operations=[
self.add_operation, self.indexing_operation,
self.unique_with_counts_operation
])
self.assertEqual(actual[1], expected_index_1)
# Equality of Value objects is done by comparing the wrapped values. Check
# the names in input_values too.
for actual_value, expected_value in zip(
actual[1].input_values, expected_index_1.input_values):
self.assertIsInstance(actual_value, value_module.InputValue)
self.assertEqual(actual_value.name, expected_value.name)
# Check that all extracted examples actually work by eval-ing them.
for example in actual:
namespace_dict = {'tf': tf}
self.assertLen(example.input_values, example.num_inputs)
for input_value in example.input_values:
namespace_dict[input_value.name] = input_value.value
eval_output = eval(example.expression, namespace_dict) # pylint: disable=eval-used
self.assertEqual(value_module.OutputValue(eval_output),
example.output_value)
self.assertEqual(example.output_value, self.example_value_1)
def test_add_constants_and_inputs_and_print_checks_names(self, bad_name):
values_by_weight = [collections.OrderedDict() for _ in range(100)]
examples = [
benchmark_module.Example(inputs={bad_name: [1, 2, 3]},
output=[1, 2])
]
benchmark = benchmark_module.Benchmark(examples, constants=[])
output_value = value_module.OutputValue(benchmark.examples[0].output)
with self.assertRaises(ValueError):
value_search._add_constants_and_inputs_and_print(
values_by_weight, benchmark, output_value,
self._constant_operation, self.settings)
def test_reconstruct_all_expressions_with_input_names_using_addition(self):
constants = [value.ConstantValue(i) for i in range(10)]
add_operation = function_operation.FunctionOperation(
tf_functions.FunctionInfo(
name='tf.add(x, y)',
filter_group=filter_group.FilterGroup.NONE,
weight=1))
# The i-th element contains all unique Value objects of weight i, mapped to
# themselves to allow retrieving the stored Value equal to some query Value.
values_by_weight = [collections.OrderedDict()] # Nothing of weight 0.
# Add constants with weight 1.
values_by_weight.append(collections.OrderedDict())
for constant in constants:
values_by_weight[1][constant] = constant
for weight in range(2, 6):
new_values = collections.OrderedDict()
for arg_1_weight in range(1, weight):
arg_2_weight = weight - arg_1_weight - 1
for arg1, arg2 in itertools.product(
values_by_weight[arg_1_weight],
values_by_weight[arg_2_weight]):
result = add_operation.apply([arg1, arg2], self.settings)
if result not in new_values:
new_values[result] = result
else:
new_values[result].merge_reconstructions(result)
values_by_weight.append(new_values)
query = value.OutputValue(9)
# The form must be (a + b), where there are 10 choices for a, which then
# determines b.
reconstructions = (values_by_weight[3][query].
reconstruct_all_expressions_with_input_names())
self.assertLen(reconstructions, 10)
# No expression uses input values.
self.assertTrue(
all(not bool(used_names) for _, used_names in reconstructions))
# No AST with only binary operators has weight 4.
self.assertEmpty(values_by_weight[4])
# The form is either (a + (b + c)) or ((a + b) + c). Each of the two forms
# has 1 + 2 + ... + 9 = 45 options. Note that "a" in (a + (b + c)) cannot be
# 0, or else (b + c) would have the same value as the entire expression.
# Similarly, "c" in ((a + b) + c) cannot be 0.
self.assertLen(
values_by_weight[5]
[query].reconstruct_all_expressions_with_input_names(), 90)
def test_normalize_names_and_extract_values(self):
name_counter = {'in3': 3, 'in4': 1, 'NEW_INPUT': 2}
copied_value = self.example_value_2.copy()
name_value_map = collect_tensor_data.normalize_names_and_extract_values(
copied_value, name_counter)
self.assertEqual(
copied_value.reconstruct_expression(use_cache=False),
# in3 -> in1, in4 -> in2, and NEW_INPUT -> in3 and in4
'tf.add(tf.add(in3, in1), tf.add(tf.add(in1, in4), tf.add(in2, in1)))'
)
self.assertEqual(
name_value_map,
# Check the raw values.
{
'in1': value_module.OutputValue([1]),
'in2': value_module.OutputValue([2]),
'in3': value_module.OutputValue([20]),
'in4': value_module.OutputValue([10])
})
for name in name_value_map:
self.assertIsInstance(name_value_map[name],
value_module.InputValue)
self.assertEqual(name, name_value_map[name].name)
def test_add_constants_and_inputs_and_print_handles_description(
self, description, expect_print, mock_stdout):
examples = [
benchmark_module.Example(inputs=[[1, 4], [2, 7]], output=[3, 11])
]
benchmark = benchmark_module.Benchmark(examples,
description=description)
values_by_weight = [collections.OrderedDict() for _ in range(5)]
output_value = value_module.OutputValue(benchmark.examples[0].output)
value_search._add_constants_and_inputs_and_print(
values_by_weight, benchmark, output_value,
self._constant_operation, self.settings)
if expect_print:
self.assertIn('Description: {}\n'.format(description),
mock_stdout.getvalue())
else:
self.assertNotIn('Description:', mock_stdout.getvalue())
def test_create_tf_examples(self):
sparse_tensor = tf.SparseTensor(values=[0, -15, 30],
indices=[[12], [34], [56]],
dense_shape=[100])
# This example does not represent a realistic tensor transformation. It uses
# variety in the input/output tensors to exercise the featurization.
io_example = collect_tensor_data.IOExample(
expression='tf.dummy_expression(in1, in2)',
input_values=[
value_module.InputValue(
[[[0.5, 2.5, 9.0], [-0.25, 0.0, 1.25]]], 'in1'),
value_module.InputValue(sparse_tensor, 'in2'),
],
output_value=value_module.OutputValue([[1.0], [0.0], [1.0],
[0.0]]),
num_inputs=2,
operations=[
self.add_operation, self.add_operation, self.gather_operation
])
with mock.patch.object(collect_tensor_data,
'COUNT_BOUNDARIES',
new=[0, 1, 3, 50, float('inf')]):
with mock.patch.object(
collect_tensor_data,
'FLOAT_BOUNDARIES',
new=[-float('inf'), -10, -1e-8, 1e-8, 10,
float('inf')]):
tf_examples = collect_tensor_data.create_tf_examples(
io_example)
operation_list = all_operations.get_operations(
include_sparse_operations=True)
expected_operations = [
2 if op.name == 'tf.add(x, y)' else
1 if op.name == 'tf.gather(params, indices)' else 0
for op in operation_list
]
expected_tf_example_1 = {
# Features from featurize_value.
'kind': [2, 2, 3, 0],
'dtype': [8, 8, 0, 0],
'rank': [2, 3, 1, 0],
'shape': [4, 1, 0, 0, 1, 2, 3, 0, 100, 0, 0, 0, 0, 0, 0, 0],
'shape_buckets': [2, 1, 0, 0, 1, 1, 2, 0, 3, 0, 0, 0, 0, 0, 0, 0],
'floats': [
1.0, 0.0, 0.5, 0.5, 9.0, -0.25, 13 / 6, 13.5 / 6, 30, -15, 5,
15, 0, 0, 0, 0
],
'float_buckets': [3, 2, 3, 3, 3, 1, 3, 3, 4, 0, 3, 4, 2, 2, 2, 2],
'counts': [
4, 4, 2, 2, 2, 0, 4, 2, 6, 6, 4, 1, 0, 1, 2, 6, 100, 3, 1, 1,
0, 1, 1, 3, 1, 1, 0, 1, 0, 0, 1, 1
],
'count_buckets': [
2, 2, 1, 1, 1, 0, 2, 1, 2, 2, 2, 1, 0, 1, 1, 2, 3, 2, 1, 1, 0,
1, 1, 2, 1, 1, 0, 1, 0, 0, 1, 1
],
'fractions': [
4 / 4, 2 / 4, 2 / 4, 2 / 4, 0 / 4, 4 / 4, 2 / 4, 6 / 6, 4 / 6,
1 / 6, 0 / 6, 1 / 6, 2 / 6, 6 / 6, 3 / 100, 1 / 3, 1 / 3,
0 / 3, 1 / 3, 1 / 3, 3 / 3, 1 / 1, 0 / 1, 1 / 1, 0 / 1, 0 / 1,
1 / 1, 1 / 1
],
'booleans': [
1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0,
0, 0, 1, 1, 1, 0, 1, 0, 1, 1, 1
],
'value_string': [
b'tf.float32:[[1.0], [0.0], [1.0], [0.0]]',
b'tf.float32:[[[0.5, 2.5, 9.0], [-0.25, 0.0, 1.25]]]',
(b'SparseTensor(indices=tf.Tensor(\n[[12]\n [34]\n [56]], '
b'shape=(3, 1), dtype=int64), '
b'values=tf.Tensor([ 0 -15 30], shape=(3,), dtype=int32), '
b'dense_shape=tf.Tensor([100], shape=(1,), dtype=int64))'),
b'0'
],
# Features from featurize_input_and_output.
'io_comparisons':
[2, 2, 2, 0, 2, 2, 1, 2, 0, 0, 2, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1],
'io_booleans': [
0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0
],
'io_counts': [1, 2, 1, 1, 2, 1, 1, 1, 1],
'io_count_buckets': [1, 1, 1, 1, 1, 1, 1, 1, 1],
'io_fractions': [
1 / 6, 2 / 4, 1 / 6, 1 / 4, 1 / 3, 2 / 4, 1 / 3, 1 / 4, 1 / 1,
1 / 1, 1 / 1, 1 / 1
],
# Features added in create_examples.
'num_inputs': [2],
'operations':
expected_operations,
'expression': [b'tf.dummy_expression(in1, in2)'],
}
print(tf_examples)
self.assertLen(tf_examples, 2)
actual_tf_example_1, actual_tf_example_2 = tf_examples
# Check the entire first example.
for key, expected in six.iteritems(expected_tf_example_1):
some_list = actual_tf_example_1.features.feature[key]
if some_list.HasField('float_list'):
actual = some_list.float_list.value
actual = [round(f, 6) for f in actual]
expected = [round(f, 6) for f in expected]
elif some_list.HasField('int64_list'):
actual = some_list.int64_list.value
elif some_list.HasField('bytes_list'):
actual = some_list.bytes_list.value
else:
self.fail('Failed to extract list from TF example.')
# Printing side-by-side like this in the test log is more helpful than the
# AssertionError message. Look at the Python3 log, which prints ints
# without the L suffix.
print('key: {}\n'
' expected: {}\n'
' got: {}'.format(key, expected, actual))
self.assertEqual(actual, expected)
# Less rigorous checks for the second example, where the two inputs have
# swapped.
self.assertEqual(
actual_tf_example_2.features.feature['rank'].int64_list.value,
[2, 1, 3, 0])
self.assertEqual(
actual_tf_example_2.features.feature['shape'].int64_list.value,
[4, 1, 0, 0, 100, 0, 0, 0, 1, 2, 3, 0, 0, 0, 0, 0])
def _find_solutions(
benchmark: benchmark_module.Benchmark,
operations: List[operation_base.Operation],
start_time: float,
settings: settings_module.Settings
) -> Tuple[List[Solution], Set[value_module.Value], ValuesByWeight,
Optional[operation_statistics.OperationStatistics]]:
"""Helper, returning (solutions, value_set, values_by_weight, statistics)."""
timeout_reached = False
end_time = start_time + settings.timeout
only_minimal_solutions = settings.only_minimal_solutions
if settings.max_solutions == 1:
# If we only want one solution, it will be minimal.
only_minimal_solutions = True
# An object to track statistics, if requested.
statistics = (operation_statistics.OperationStatistics()
if settings.printing.statistics
else None)
# A list of Solution namedtuples.
solutions = []
# A set of string solution expressions (don't return duplicate solutions).
solution_expression_set = set()
# The output value to search for.
output_value = value_module.OutputValue(benchmark.examples[0].output)
# A list of OrderedDicts mapping Value objects to themselves. The i-th
# OrderedDict contains all Value objects of weight i.
values_by_weight = [collections.OrderedDict()
for _ in range(settings.max_weight + 1)]
# Find and cache the cast and constant operations for use later.
cast_operation = None
constant_operation = None
for operation in operations:
if operation.name == tf_functions.CAST_OPERATION_NAME:
cast_operation = operation
elif operation.name == tf_functions.CONSTANT_OPERATION_NAME:
constant_operation = operation
# Create the output dtype value for use later.
dtype_value = value_module.ConstantValue(output_value.dtype)
# Populate values_by_weight with inputs and constants. This also prints
# inputs/output/constants to stdout.
_add_constants_and_inputs_and_print(
values_by_weight, benchmark, output_value, constant_operation, settings)
# A set storing all values found so far.
value_set = set().union(*values_by_weight)
filter_cache = filtered_values_cache.FilteredValuesCache()
# Value search by weight.
for weight in range(1, settings.max_weight + 1):
if settings.printing.progress:
print('Searching weight {}...'.format(weight))
# Values with the current weight. This might already include leaf values.
new_values = values_by_weight[weight]
for operation in operations:
for value in operation.enumerate_values_with_weight(
target_weight=weight,
values_by_weight=values_by_weight,
filter_cache=filter_cache,
end_time=end_time,
settings=settings,
statistics=statistics):
if value not in value_set:
# This value has never been seen before, or it's the desired output.
if settings.printing.verbose:
expression = value.reconstruct_expression()
print('{} produces:\n{}'.format(expression, value))
if value == output_value:
possible_first_solution = not solutions
# Found solution(s), but some may be bad.
_record_solutions(value, weight, start_time, solutions,
solution_expression_set, benchmark, settings)
if possible_first_solution and solutions:
end_time = min(
end_time,
timeit.default_timer() + settings.max_extra_solutions_time)
if len(solutions) >= settings.max_solutions:
return solutions, value_set, values_by_weight, statistics
else:
# Only store the value if it isn't a solution. Otherwise, we'll get
# lots of "almost duplicate" solutions, e.g., by adding 0.
new_values[value] = value
# We should never add output_value (or anything equal) to value_set
# so that we can continue finding other solutions.
value_set.add(value)
else: # This value has been seen before.
if value in new_values:
# The value was already computed differently with this weight.
original_value = new_values[value]
if isinstance(original_value, value_module.OperationValue):
# Only merge reconstructions if this was originally an
# OperationValue. (It could be a ConstantValue instead.)
operation_value = original_value # type: value_module.OperationValue
operation_value.merge_reconstructions(value)
elif not only_minimal_solutions:
# If we want non-minimal solutions, we need to store the value even
# if we have already seen that value with a smaller weight.
new_values[value] = value
if timeit.default_timer() > end_time:
timeout_reached = True
# Don't return immediately; still try to cast new values because this is
# relatively quick.
break
# Try casting new values to the output dtype if this has a chance of being
# a correct solution.
for new_value in new_values:
if (cast_operation is not None and
new_value.shape == output_value.shape and
new_value.dtype != output_value.dtype and
operation_filtering.is_castable(new_value, dtype_value)):
casted_value = cast_operation.apply([new_value, dtype_value], settings)
if casted_value == output_value:
possible_first_solution = not solutions
# Found solution(s), but some may be bad.
_record_solutions(casted_value, weight, start_time, solutions,
solution_expression_set, benchmark, settings)
if possible_first_solution and solutions:
end_time = min(
end_time,
timeit.default_timer() + settings.max_extra_solutions_time)
if len(solutions) >= settings.max_solutions:
return solutions, value_set, values_by_weight, statistics
if settings.printing.progress:
print('Found {} distinct values of weight {}, or {} total.'.format(
len(new_values), weight, len(value_set)))
if only_minimal_solutions and solutions:
return solutions, value_set, values_by_weight, statistics
if timeout_reached:
break
return solutions, value_set, values_by_weight, statistics
def test_reconstruct_expression(self):
output_value = value.OutputValue(1)
with self.assertRaises(NotImplementedError):
output_value.reconstruct_expression()
def test_copy(self):
with self.assertRaises(NotImplementedError):
value.OutputValue(1).copy()
def test_add_constants_and_inputs_and_print(self, mock_stdout):
values_by_weight = [collections.OrderedDict() for _ in range(100)]
examples = [
benchmark_module.Example(
inputs=[
[[[[1, 2, 3, 4, 5, 6, 7, 8, 9]]]], # Shape (1, 1, 1, 9).
[list(range(6)) for _ in range(15)], # Shape (15, 6).
42,
],
output=tf.constant(list(range(50)),
dtype=tf.int64), # Shape (50).
),
]
constants = [123, 45, 1, 9]
benchmark = benchmark_module.Benchmark(examples, constants)
output_value = value_module.OutputValue(benchmark.examples[0].output)
value_search._add_constants_and_inputs_and_print(
values_by_weight, benchmark, output_value,
self._constant_operation, self.settings)
expected_constants_ordered = [
# Provided constants.
_constant(123),
_constant(45),
_constant(1),
_constant(9),
# Input variables.
value_module.InputValue([1], 'in1'), # Reconstructs to 'in1'.
value_module.InputValue([2], 'in2'), # Reconstructs to 'in2'.
value_module.InputValue([3], 'in3'), # Reconstructs to 'in3'.
# Common constants.
_constant(0),
_constant(-1),
_constant(True),
_constant(False),
# Common DTypes.
_constant(tf.int32),
_constant(tf.float32),
_constant(tf.bool),
# Primitive inputs as scalar tensors. The actual Value is an
# OperationValue, but this ExpressionValue has the same expression.
value_module.ExpressionValue(None, 'tf.constant(in3)'),
# Axis constants.
_constant(2),
_constant(3),
# Uncomon DTypes.
_constant(tf.int64),
# Shape constants.
_constant(6),
_constant(15),
_constant(50),
# Output shape tuple.
_constant((50, )),
]
# values_by_weight is a list of OrderedDicts. Convert to list of lists.
actual_constants_ordered = sum(
[list(ordered_dict) for ordered_dict in values_by_weight], [])
self.assertEqual(
[v.reconstruct_expression() for v in actual_constants_ordered],
[v.reconstruct_expression() for v in expected_constants_ordered])
# The output shape tuple (50,) is included as a constant, but not printed.
self.assertIn(
'Constants: [123, 45, 1, 9, 0, -1, True, False, 2, 3, 6, 15, 50]\n',
mock_stdout.getvalue())
def test_hash(self): value_1a = value.InputValue(tf.constant([1, 2]), 'a') value_1b = value.OutputValue(tf.constant([1, 2])) value_2 = value.InputValue(tf.constant([1, 3]), 'a') self.assertEqual(hash(value_1a), hash(value_1b)) self.assertNotEqual(hash(value_1a), hash(value_2))
def test_apply(self):
self.assertEqual(
self.operation.apply(self.arg_values,
settings_module.default_settings()),
value.OutputValue([34, 78]))
def compute_search_space_size(benchmark, settings, description_handler):
"""Computes and prints the size of the search space.
This counts the total number of expressions with weight at most max_weight.
The weights come from the benchmark (for constants and inputs) and the
description handler (for determining the op weights). Distinct expressions
will be counted separately even if they evaluate to the same value, unlike in
TF-Coder's value_search algorithm which does value-based pruning.
Args:
benchmark: The Benchmark object defining the problem to analyze.
settings: A Settings object containing settings for value search.
description_handler: The DescriptionHandler used, which can modify weights
of operations.
Returns:
Nothing. All output is printed to stdout.
"""
max_weight = settings.max_weight
print('Computing search space.\n'
'Benchmark name: {}\n'
'Description handler: {}\n'
'Max weight: {}'.format(benchmark.name, description_handler,
max_weight))
# TODO(kshi): Update to load the tensor features model/config.
operations = value_search.get_reweighted_operations(benchmark,
settings,
description_handler,
tensor_model=None,
tensor_config=None)
# These loops are not the most efficient, but it doesn't really matter.
print('\nFound {} operations.'.format(len(operations)))
print()
for weight in range(1, max(op.weight for op in operations) + 1):
print('# operations with weight {}: {}'.format(
weight, sum(1 for op in operations if op.weight == weight)))
print()
for arity in range(1, max(op.num_args for op in operations) + 1):
print('# operations with arity {}: {}'.format(
arity, sum(1 for op in operations if op.num_args == arity)))
output_value = value_module.OutputValue(benchmark.examples[0].output)
values_by_weight = [
collections.OrderedDict() for _ in range(max_weight + 1)
]
constant_operation = None
for operation in operations:
if operation.name == tf_functions.CONSTANT_OPERATION_NAME:
constant_operation = operation
break
with SuppressPrint():
value_search._add_constants_and_inputs_and_print( # pylint: disable=protected-access
values_by_weight, benchmark, output_value, constant_operation,
settings)
num_expressions_with_weight = [
len(values_with_weight) for values_with_weight in values_by_weight
]
print()
max_weight_with_initial_value = max(w for w in range(max_weight + 1)
if num_expressions_with_weight[w])
for weight in range(1, max_weight_with_initial_value + 1):
print('# initial values with weight {}: {}'.format(
weight, num_expressions_with_weight[weight]))
for total_weight in range(2, max_weight + 1):
for operation in operations:
# All operations should have strictly positive weight and num_args.
op_weight = operation.weight
op_arity = operation.num_args
if total_weight - op_weight < op_arity:
continue
# Partition `total_weight - op_weight` into `op_arity` positive pieces.
# Equivalently, partition `total_weight - op_weight - op_arity` into
# `op_arity` nonnegative pieces.
for partition in tf_coder_utils.generate_partitions(
total_weight - op_weight - op_arity, op_arity):
arg_weights = [part + 1 for part in partition]
num_expressions_with_weight[total_weight] += functools.reduce(
operator.mul,
(num_expressions_with_weight[w] for w in arg_weights))
print()
for weight in range(1, max_weight + 1):
print('# expressions with weight exactly {}: {}'.format(
weight, num_expressions_with_weight[weight]))
print()
for weight in range(1, max_weight + 1):
print('# expressions with weight up to {}: {}'.format(
weight, sum(num_expressions_with_weight[:weight + 1])))
def test_init_converts_to_tensor(self): output_value = value.OutputValue(1) self.assertTrue(output_value.is_tensor)
