def test_get_operations_consistent_order(self):
operations_1 = all_operations.get_operations(
include_sparse_operations=True)
operations_2 = all_operations.get_operations(
include_sparse_operations=True)
self.assertEqual([op.name for op in operations_1],
[op.name for op in operations_2])
def test_get_operations_correct_type(self):
operations = all_operations.get_operations(
include_sparse_operations=True)
self.assertTrue(
all(
isinstance(element, operation_base.Operation)
for element in operations))
def test_get_operations_includes_expected(self, name, is_sparse):
for include_sparse_operations in [True, False]:
operations = all_operations.get_operations(
include_sparse_operations=include_sparse_operations)
should_be_included = include_sparse_operations or not is_sparse
self.assertEqual(
any(operation.name == name for operation in operations),
should_be_included)
def create_examples(io_example: IOExample,
max_num_inputs: int = 3,
permute_inputs: bool = True) -> List[Dict[Text, Any]]:
"""Creates example dicts for the I/O example."""
examples = []
operation_list = all_operations.get_operations(
include_sparse_operations=True)
operation_counter = collections.Counter(
[op.name for op in io_example.operations])
operation_counts = [operation_counter[op.name] for op in operation_list]
num_inputs = len(io_example.input_values)
try:
num_inputs_feature = min(num_inputs, max_num_inputs)
output_features = featurize_value(io_example.output_value)
input_features = []
for input_value in io_example.input_values[:max_num_inputs]:
combined_features = featurize_value(input_value)
combined_features.update(
featurize_input_and_output(input_value,
io_example.output_value))
input_features.append(combined_features)
dummy_value = value_module.ConstantValue(0)
dummy_input_features = featurize_value(dummy_value)
dummy_input_features.update(
featurize_input_and_output(dummy_value, dummy_value))
except ValueError as e:
logging.warning('%s: could not featurize IOExample %s', e, io_example)
return []
permutations = (itertools.permutations(range(num_inputs))
if permute_inputs else [list(range(num_inputs))])
for permutation in permutations:
feature_dict = collections.defaultdict(list)
feature_dict['num_inputs'] = [num_inputs_feature]
feature_dict.update(copy.deepcopy(output_features))
padded_input_features = [
input_features[index] for index in permutation
]
for _ in range(max_num_inputs - num_inputs):
padded_input_features.append(dummy_input_features)
for input_features_dict in padded_input_features:
for key, value in six.iteritems(input_features_dict):
feature_dict[key].extend(value)
feature_dict['operations'] = operation_counts
feature_dict['expression'] = [io_example.expression]
examples.append(feature_dict)
return examples
def setUp(self):
super(ValueSearchTest, self).setUp()
operations = all_operations.get_operations()
self._constant_operation = None
for operation in operations:
if operation.name == tf_functions.CONSTANT_OPERATION_NAME:
self._constant_operation = operation
break
self.assertIsNotNone(self._constant_operation)
self.settings = settings_module.from_dict({'timeout': 10})
def __init__(self,
operations: Optional[List[operation_base.Operation]] = None):
"""Initializes the handler.
Args:
operations: A list of operations that the scorer should handle. Exposed
for testing.
Raises:
ValueError: If there are duplicate operation names.
"""
self.operations = (operations
if operations else all_operations.get_operations(
include_sparse_operations=True))
self.all_names = [operation.name for operation in self.operations]
if len(set(self.all_names)) != len(self.operations):
raise ValueError('Duplicate operation name.')
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 get_reweighted_operations(
benchmark: benchmark_module.Benchmark,
settings: settings_module.Settings,
description_handler: Optional[DescriptionHandler] = None,
tensor_model: Optional[tensor_features_model.Model] = None,
tensor_config: Optional[Dict[Text, Any]] = None,
) -> List[operation_base.Operation]:
"""Returns a list of operations with correct weights for the problem."""
include_sparse_operations = (
not settings.operations.limit_sparse_operations
or _contains_sparse(benchmark))
operations = all_operations.get_operations(
include_sparse_operations=include_sparse_operations)
operation_names = [op.name for op in operations]
if len(operation_names) != len(set(operation_names)):
raise ValueError('Operation names were not unique.')
if settings.paper_experiments.uniform_weights:
# Only for experiments in the PLDI paper.
for operation in operations:
operation.weight = 1
return operations
multipliers = {}
if description_handler and benchmark.description:
multipliers = _combine_multipliers(
multipliers,
description_handler.get_operation_multipliers(
benchmark.description, settings))
if tensor_model is not None and tensor_config is not None:
multipliers = _combine_multipliers(
multipliers,
operation_multipliers_from_tensor_model(benchmark, tensor_model,
tensor_config, settings))
for operation in operations:
operation.weight = max(
1,
int(round(operation.weight * multipliers.get(operation.name, 1))))
return operations
def get_config():
"""Gets the default config for the tensor features model."""
cfg = {}
# Model size.
cfg['feedforward_size'] = 2048
cfg['num_feedforward_layers'] = 2
# Training hyperparams.
cfg['batch_size'] = 1024
cfg['learning_rate'] = 1e-5 # 1e-5 is better, 1e-4 is faster for debugging.
cfg['global_norm_clip'] = 5
cfg['num_epochs'] = 3
# Loss function and options.
# sigmoid_ce: Average the sigmoid cross entropy loss over all operations.
# f_beta: Use a differentiable version of F_beta score as the loss.
cfg['loss'] = 'f_beta'
# If using f_beta loss, this is the value of beta. Because TF-Coder cares more
# about recall than precision, this value should be >= 1.
cfg['beta'] = 2
# Whether to use weights to "balance" the prevalence of different ops. Let
# `c_i` be the number of training examples operation `op_i` appears in. If
# this option is True, then if `op_i` is used in a given example, we treat it
# as actually appearing ((mean_j c_j) / c_i) times. This effectively upweights
# rare operations and downweights common operations, such that each operation
# now appears an equal (mean_j c_j) times across the dataset. Furthermore,
# the total number of operations used across the dataset remains the same,
# namely (sum_j c_j).
cfg['weighted_ops'] = True
# If 'mean', use (mean_j c_j) as the weight numerator. If 'max', use
# (max_j c_j) instead, so that no operations are downweighted.
cfg['weight_numerator'] = 'mean'
# If using weights to balance the occurrences of ops, this is the maximum
# weight. Weights that are too large may increase training instability.
# However, we do have gradient clipping in place, so this may be unnecessary.
cfg['max_weight'] = 10000
# Details about the training data.
operations = all_operations.get_operations(include_sparse_operations=True)
cfg['num_ops'] = len(operations)
cfg['num_kinds'] = 5
cfg['num_dtypes'] = len(tf_coder_utils.INT_DTYPES +
tf_coder_utils.FLOAT_DTYPES + (tf.bool, ))
cfg['max_rank'] = collect_tensor_data.MAX_RANK
cfg['num_shape_buckets'] = len(collect_tensor_data.COUNT_BOUNDARIES)
cfg['num_float_buckets'] = len(collect_tensor_data.FLOAT_BOUNDARIES)
cfg['num_count_buckets'] = len(collect_tensor_data.COUNT_BOUNDARIES)
cfg['max_num_inputs'] = 3
# Store the names of all operations used at the time of model training, so we
# know how to interpret the model's predictions at test time, even if new
# operations were added in between.
cfg['operation_names'] = [op.name for op in operations]
cfg['op_counts'] = [
261791, 3219730, 49984, 468220, 70588, 671391, 304105, 76203, 361,
12293742, 223803, 208334, 4469035, 0, 5228323, 395630, 514399, 8810328,
175664, 982, 48877, 24769, 1429728, 15448, 124548, 675, 788727, 180186,
770114, 198687, 7687, 26592, 773463, 443414, 26209, 183343, 397296,
580331, 257051, 22749, 24187, 1229, 1727, 8011, 28719, 316277, 229433,
54625, 5881, 23486, 2548, 33794, 466112, 1906853, 590050, 2780775,
29209, 2393112, 4041, 342637, 39436, 9667, 1308, 2748, 752650, 17687,
5663, 1246906, 1327261, 486676, 421095, 126426, 184234, 134355,
2565768, 2124113, 5845, 49484, 338106, 63409, 6227, 9410, 953093,
124307, 878241, 126966, 211238, 96981, 185479, 4384591, 800148, 119813,
188229, 5031788, 1761827, 244750, 2493965, 2, 339585, 1145115, 4293436,
149049, 1058, 342232, 687, 33192, 451, 31190, 1145201, 2621, 1642,
128519, 142314, 22508, 16, 3660, 499, 48583, 359399, 118306, 3940,
9318, 3614, 543957, 809717, 9070582, 2120948, 273953, 719889, 603656,
170092, 517624, 481419, 2997663
]
# Evaluation settings.
cfg['eval_batch_size'] = 100000 # This is actually the entire eval set.
cfg['eval_step_frequency'] = 100 # A multiple of summary_step_frequency.
# Saving summaries and checkpoints.
cfg['summary_step_frequency'] = 100
cfg['save_step_frequency'] = 100
cfg['keep_checkpoint_every_n_hours'] = 0.5
cfg['allow_restore'] = False
cfg['log_device_placement'] = False
return cfg
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 compute_search_space_size(num_ops, num_nodes, num_leaf_choices):
"""Computes and prints the size of the search space.
This counts the total number of expressions with exactly the given number of
operations and nodes in the expression tree. 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:
num_ops: The target number of operations.
num_nodes: A target number of nodes in the expression tree.
num_leaf_choices: The number of distinct inputs and constants available to
form the leaves of the expression tree.
Returns:
The DP table, where dp[i][j] is the answer for i ops and j nodes.
"""
operations = all_operations.get_operations(include_sparse_operations=True)
max_arity = max(op.num_args for op in operations)
arity_counts = [0] * (max_arity + 1)
print('Found {} operations.'.format(len(operations)))
for arity in range(max_arity + 1):
arity_counts[arity] = sum(1 for op in operations
if op.num_args == arity)
print('# operations with arity {}: {}'.format(arity,
arity_counts[arity]))
print('\nNum leaf nodes: {}'.format(num_leaf_choices))
# dp[i][j] = the number of expressions using exactly i ops and j nodes.
dp = np.zeros((num_ops + 1, num_nodes + 1))
# The only expressions using 0 ops are single-node leaves.
dp[0][1] = num_leaf_choices
for ops in range(1, num_ops + 1):
for nodes in range(1, num_nodes + 1):
# The running total number of ways to satisfy # ops and # nodes.
total = 0
for arity in range(1, max_arity + 1):
# The running total number of ways to fill the arguments.
args_total = 0
# The ways to allocate remaining ops and nodes to the arguments.
ops_partitions = utils.generate_partitions(num_elements=ops -
1,
num_parts=arity)
nodes_partitions = utils.generate_partitions(
num_elements=nodes - 1, num_parts=arity)
for ops_partition, nodes_partition in itertools.product(
ops_partitions, nodes_partitions):
# The i-th argument must have ops_partition[i] ops and
# nodes_partition[i] nodes. Look up the number of ways in the DP
# table.
args_total += np.prod([
dp[ops_partition[i]][nodes_partition[i]]
for i in range(arity)
])
# There are arity_counts[arity] choices for the outermost operation.
total += args_total * arity_counts[arity]
dp[ops][nodes] = total
print(
'\nThere are {} expressions using exactly {} ops and {} nodes.'.format(
dp[num_ops][num_nodes], num_ops, num_nodes))
print('There are {} expressions using at most {} ops and {} nodes.'.format(
np.sum(dp), num_ops, num_nodes))
return dp
def test_get_operations_all_have_docstrings(self):
operations = all_operations.get_operations(
include_sparse_operations=True)
self.assertTrue(
all(operation.metadata.docstring for operation in operations))
def test_get_operations_unique_names(self):
operations = all_operations.get_operations(
include_sparse_operations=True)
names_set = {operation.name for operation in operations}
self.assertLen(names_set, len(operations))
def test_get_operations_correct_cardinality(self,
include_sparse_operations,
expected_cardinality):
operations = all_operations.get_operations(
include_sparse_operations=include_sparse_operations)
self.assertLen(operations, expected_cardinality)
def test_value_search_prioritizes_operations(self):
# Make sure prioritized ops get printed out.
settings = settings_module.from_dict({
'timeout':
0.5, # We don't actually care about solving the problem.
'printing.prioritized_operations':
True,
'printing.deprioritized_operations':
True,
'tensor_model.prioritize_threshold':
0.5,
'tensor_model.deprioritize_threshold':
0.2,
})
# Create data for the mocks.
all_ops = all_operations.get_operations(include_sparse_operations=True)
op_names = [op.name for op in all_ops]
abs_index = None # Prioritized by tensor features model.
argmax_index = None # Deprioritized by tensor features model.
for i, op in enumerate(all_ops):
if op.name == 'tf.abs(x)':
abs_index = i
elif op.name == 'tf.argmax(input, axis)':
argmax_index = i
self.assertIsNotNone(abs_index)
self.assertIsNotNone(argmax_index)
operation_probs = np.repeat(0.4, len(all_ops))
operation_probs[abs_index] = 0.6 # Above prioritize threshold.
operation_probs[argmax_index] = 0.1 # Below deprioritize threshold.
operation_logits = np.expand_dims(_inverse_sigmoid(operation_probs),
axis=0)
nl_data = [
# Associate "apple" with tf.zeros.
{
'docstring': ['apple pie'],
'tf_functions': ['tf.zeros'],
'comments': [],
'names': [],
'strings': []
},
] * 100
# Mock the tensor model's predictions, the NL model's data, and print().
with mock.patch('tf_coder.models.tensor_features_model.eval_single_example',
return_value=tensor_features_model.Result(
operation_logits=operation_logits)), \
mock.patch('tf_coder.datasets.github.data_loader.load_data',
return_value=nl_data), \
mock.patch('sys.stdout', new_callable=six.StringIO) as mock_stdout:
handler = bag_of_words_handlers.NaiveBayesDescriptionHandler(
max_num_prioritized=1)
benchmark = benchmark_module.Benchmark(
# I/O example doesn't matter.
examples=[
benchmark_module.Example({'my_var': [1, 2]}, [2, 1])
],
# Description contains "apple"!
description='honeycrisp apple')
value_search.run_value_search(benchmark=benchmark,
description_handler=handler,
settings=settings,
tensor_model=mock.Mock(),
tensor_config={
'max_num_inputs': 3,
'operation_names': op_names
})
self.assertIn('BOW handler prioritized tf.zeros',
mock_stdout.getvalue())
self.assertIn('Tensor features model prioritized tf.abs(x), p=0.6',
mock_stdout.getvalue())
self.assertIn(
'Tensor features model deprioritized tf.argmax(input, axis), p=0.1',
mock_stdout.getvalue())
