def encode(state, value):
"""Encode tf_computation."""
encoded_structure = nest_contrib.map_structure_up_to(
encoders, lambda state, value, e: e.encode(value, state), state,
value, encoders)
encoded_value = nest_contrib.map_structure_up_to(
encoders, lambda s: s[0], encoded_structure)
new_state = nest_contrib.map_structure_up_to(encoders, lambda s: s[1],
encoded_structure)
return new_state, encoded_value
def encode_fn(x, encode_params, decode_before_sum_params):
encoded_structure = nest_contrib.map_structure_up_to(
encoders, lambda e, *args: e.encode(*args), encoders, x,
encode_params)
encoded_x = _slice(encoders, encoded_structure, 0)
state_update_tensors = _slice(encoders, encoded_structure, 1)
return encoded_x, decode_before_sum_params, state_update_tensors
def testMapStructureUpTo(self):
# Named tuples.
ab_tuple = collections.namedtuple("ab_tuple", "a, b")
op_tuple = collections.namedtuple("op_tuple", "add, mul")
inp_val = ab_tuple(a=2, b=3)
inp_ops = ab_tuple(a=op_tuple(add=1, mul=2), b=op_tuple(add=2, mul=3))
out = nest.map_structure_up_to(
inp_val, lambda val, ops: (val + ops.add) * ops.mul, inp_val, inp_ops)
self.assertEqual(out.a, 6)
self.assertEqual(out.b, 15)
# Lists.
data_list = [[2, 4, 6, 8], [[1, 3, 5, 7, 9], [3, 5, 7]]]
name_list = ["evens", ["odds", "primes"]]
out = nest.map_structure_up_to(
name_list, lambda name, sec: "first_{}_{}".format(len(sec), name),
name_list, data_list)
self.assertEqual(out, ["first_4_evens", ["first_5_odds", "first_3_primes"]])
# Dicts.
inp_val = dict(a=2, b=3)
inp_ops = dict(a=dict(add=1, mul=2), b=dict(add=2, mul=3))
out = nest.map_structure_up_to(
inp_val,
lambda val, ops: (val + ops["add"]) * ops["mul"], inp_val, inp_ops)
self.assertEqual(out["a"], 6)
self.assertEqual(out["b"], 15)
# Non-equal dicts.
inp_val = dict(a=2, b=3)
inp_ops = dict(a=dict(add=1, mul=2), c=dict(add=2, mul=3))
with self.assertRaisesWithLiteralMatch(
ValueError,
nest._SHALLOW_TREE_HAS_INVALID_KEYS.format(["b"])):
nest.map_structure_up_to(
inp_val,
lambda val, ops: (val + ops["add"]) * ops["mul"], inp_val, inp_ops)
# Dict+custom mapping.
inp_val = dict(a=2, b=3)
inp_ops = _CustomMapping(a=dict(add=1, mul=2), b=dict(add=2, mul=3))
out = nest.map_structure_up_to(
inp_val,
lambda val, ops: (val + ops["add"]) * ops["mul"], inp_val, inp_ops)
self.assertEqual(out["a"], 6)
self.assertEqual(out["b"], 15)
# Non-equal dict/mapping.
inp_val = dict(a=2, b=3)
inp_ops = _CustomMapping(a=dict(add=1, mul=2), c=dict(add=2, mul=3))
with self.assertRaisesWithLiteralMatch(
ValueError,
nest._SHALLOW_TREE_HAS_INVALID_KEYS.format(["b"])):
nest.map_structure_up_to(
inp_val,
lambda val, ops: (val + ops["add"]) * ops["mul"], inp_val, inp_ops)
def get_params_fn(state):
params = nest_contrib.map_structure_up_to(encoders,
lambda e, s: e.get_params(s),
encoders, state)
encode_params = _slice(encoders, params, 0)
decode_before_sum_params = _slice(encoders, params, 1)
decode_after_sum_params = _slice(encoders, params, 2)
return encode_params, decode_before_sum_params, decode_after_sum_params
def _slice(encoders, nested_value, idx):
"""Takes a slice of nested values.
We use a collection of encoders to encode a collection of values. When a
method of the encoder returns a tuple, e.g., encode / decode params of the
get_params method, we need to recover the matching collection of encode params
and collection of decode params. This method is a utility to achieve this.
Args:
encoders: A collection of encoders.
nested_value: A collection of indexable values of the same structure as
`encoders`.
idx: An integer. Index of the values in `nested_value` along which to take
the slice.
Returns:
A collection of values of the same structure as `encoders`.
"""
return nest_contrib.map_structure_up_to(encoders, lambda t: t[idx],
nested_value)
def update_state_fn(state, state_update_tensors):
return nest_contrib.map_structure_up_to(
encoders, lambda e, *args: e.update_state(*args), encoders, state,
state_update_tensors)
def decode_after_sum_fn(summed_values, decode_after_sum_params):
part_decoded_aggregated_x, num_summands = summed_values
return nest_contrib.map_structure_up_to(
encoders,
lambda e, x, params: e.decode_after_sum(x, params, num_summands),
encoders, part_decoded_aggregated_x, decode_after_sum_params)
def decode_before_sum_tf_function(encoded_x, decode_before_sum_params):
part_decoded_x = nest_contrib.map_structure_up_to(
encoders, lambda e, *args: e.decode_before_sum(*args), encoders,
encoded_x, decode_before_sum_params)
one = tf.constant((1, ), tf.int32)
return part_decoded_x, one
def decode(encoded_value):
"""Decode tf_computation."""
return nest_contrib.map_structure_up_to(encoders,
lambda e, val: e.decode(val),
encoders, encoded_value)