Differently from the previous 2 articles, where I merged the description of the solutions of two problems into one article, this time the whole article is dedicated to the pure TensorFlow solution to problem number 5. The reason is simple: solving this problem in pure TensorFlow hasn’t been straightforward so it is worth explaining all the limitations and the subtleties found during the solution.
In the first part of the article, I’ll explain the solution that solves completely both parts of the puzzle. As usual, focusing on the TensorFlow features used during the solution and all the various technical details worth explaining. In the second part, instead, I’ll propose a potential alternative solution to the problem that uses a
tf.Variable with an “undefined shape”. This is a feature of
tf.Variable that’s not clearly documented and, thus, widely used. So, at the end of this article, we’ll understand how to solve the day 5 problem in pure TensorFlow and also have an idea of how to re-design the solution using a
tf.Variable with the
validate_shape argument set to
You can click on the title above to read the full text of the puzzle. The TLDR version is: we have an initial configuration of crates (the puzzle input) and a set of moves to perform. Part 1 constrains the crane that’s moving the crates to pick a single crate at a time, while part 2 allows multiple crates to be picked up at the same time. The problem asks us to determine, after having moved the crates, what crate ends up on top of each stack.
So, given a puzzle input like
[D] [N] [C] [Z] [M] [P] 1 2 3 move 1 from 2 to 1 move 3 from 1 to 3 move 2 from 2 to 1 move 1 from 1 to 2
Part 1, which wants us to move the crates one at a time, ends up with this final configuration
[Z] [N] [D] [C] [M] [P] 1 2 3
Thus, the result is: “CMZ”.
For part 2 instead, where multiple crates can be picked up at the same time, the final configuration is
[D] [N] [Z] [M] [C] [P] 1 2 3
in this case the result is “MCD”.
The problem can be breakdown into 4 simple steps:
- Read the first part of the input: parse the crates
- Create a data structure that models the stacks of crates
- Read the second part of the input: parse the instructions
- Iteratively transform the previously created data structure according to the instructions
Parsing strings: tf.strings & regex
tf.strings module contains several utilities for working with
dtype=tf.string. Unfortunately, there are tons of limitations when working with strings and
tf.function-decorated functions (we’ll see some of these limitations later). Moreover, perhaps the most powerful tool for string parsing and manipulation (the regular expressions) has a very limited integration. We only have 2 functions to use:
- tf.strings.regex_full_match: check if the input matches the regex pattern.
- tf.strings.regex_replace: replace elements of input matching regex pattern with rewrite.
That’s all. So, if we want to use regular expressions for parsing the crates or the moves, we are forced to use only these 2 functions, together with the other basic string manipulation functions offered by the module.
Parsing the crates
We are quite lucky because every line containing at least a crate contains the character
[. Thus, a regex pattern for this pattern is
We can thus easily filter every line that contains a crate by applying this condition to every line of the input dataset.
stacks_dataset = dataset.filter( lambda line: tf.strings.regex_full_match(line, r".*\[.*") )
Now, we need to find a way to extract every stack of crates. This can be quite easily done noticing that every stack is just a column of 4 characters. Thus, by moving 4 characters at a time over every line, we can understand where a crate is and extract its letter.
Of course, the tool to use is
tf.TensorArray since every stack could contain a variable number of crates, and
tf.TensorArray is the only tool we can use in static-graph mode for having such behavior.
@tf.function def to_array(line): length = tf.strings.length(line) + 1 stacks = length // 4 ta = tf.TensorArray(tf.string, size=0, dynamic_size=True) for i in tf.range(stacks): substr = tf.strings.strip(tf.strings.substr(line, i * 4, 4)) stripped = tf.strings.regex_replace(substr, r"\[|\]", "") ta = ta.write(i, stripped) return ta.stack()
to_array function accepts a line
[Z] [M] [P] as input, parses it by moving 4 characters at a time along its length, removes the square brackets, and saves the letter in the correct position of the stack (0 indexed).
stacks_dataset = stacks_dataset.map(to_array)
The lines are read from the first to the last, that means every stack is built from the top to the bottom.
Now that we have a dataset that produces stacks when iterated, we can convert it to a
tf.Tensor. In this way our
stacks_tensor contains a bunch of very important information:
- The number of stacks (the shape along the 1-dimension)
- The height of the highest stack
stacks_tensor = tf.convert_to_tensor(list(stacks_dataset)) num_stacks = tf.shape(stacks_tensor, tf.int64) + 1
stacks_tensor is unmodifiable not being a
tf.Variable. However, we are in a particular situation in which TensorFlow is not perfectly suited. In fact, even if
stacks_tensor' is converted into a tf.Variable` we have no idea about the maximum height that can be reached when applying the various instructions and moving the crates.
Stacking crates: the limitations of tf.Variable
As it’s easy to understand, we have no idea how tall a stack can become, everything depends on the instructions. Since a
tf.Variable requires a well-defined shape, we are forced to guess this maximum size and hope it’s enough. That’s for sure a huge limitation of this solution.
However, in the last part of the article I’ll suggest to the reader a potential starting point for designing a more general solution that uses the
validate_shape=False parameter of the
tf.Variable - that’s a not widely used feature, but very powerful.
max_stack_size = 200 stacks = tf.Variable(tf.zeros((max_stack_size, num_stacks - 1, 1), dtype=tf.string))
stacks is a
tf.Variable with shape
(200, number of stacks -1, 1):
- 200 is the maximum number of crates in a stack our solution supports (the limitation)
num_stacks -1is the number of stacks (0-based).
1is the dimension of the inner dimension. Let’s say that’s not required (since the unary dimensions can be squeezed) but having it simplifies the usage of the
stacks variable is initialized with zeros, but we have our
stacks_tensor that represents the initial state. So, we can define a function that initializes to the initial state the
stacks variable every time is called.
def initialize_stacks(): stacks.assign(tf.zeros_like(stacks)) indices_x, indices_y = tf.meshgrid( tf.range(max_stack_size - tf.shape(stacks_tensor), max_stack_size), tf.range(tf.shape(stacks_tensor)), ) indices = tf.stack([indices_x, indices_y], axis=-1) updates = tf.expand_dims(tf.transpose(stacks_tensor), axis=2) stacks.assign(tf.tensor_scatter_nd_update(stacks, indices, updates)) initialize_stacks()
First, we reset the status of the
stacks variable to zero. Then, we create the indices we want to use while assigning elements of
stacks. The idea is to assign the elements “at the bottom” which means that a stack in
ZND up to down) is assigned to its corresponding stack in the
stacks tensor at the position
max_stack_size -3 onwards (e.g.
Z in position
N in position max_stack_size-2`, and so on).
Then, we use the
tf.tensor_scatter_nd_update function to scatter
stacks according to
indices. This function returns a new tensor with the same shape of
stacks, that thus can be perfectly assigned to the
Having a single
tf.Variable containing all the stacks at once, we have lost the information about the number of crates per stack. In fact,
tf.shape(stacks) will always return the whole variable shape without information about the non-empty elements present along every dimension.
Thus, we need to keep track manually of this information, which needs to be re-computed every time we apply a transformation to the
Counting crates: mutable hashmaps
The idea is simple: the zero element for the
tf.string dtype is the empty string (
""). Thus, we can count the number of non-empty elements along the 0-axis, and save the result inside a mutable hashmap that maps the stack number to its number of crates.
Once again, the perfect tool (although still experimental, it works fine!) is tf.lookup.experimental.MutableHashTable.
num_elements = tf.lookup.experimental.MutableHashTable( tf.int64, tf.int64, default_value=-1 ) def update_num_elements(): num_elements.insert( tf.range(num_stacks - 1), tf.squeeze( tf.reduce_sum(tf.cast(tf.not_equal(stacks, ""), tf.int64), axis=) ), ) update_num_elements()
Here we go, we have the
stacks variable populated with the initial state, and the
num_elements lookup table that contains the number of elements for each stack.
We can now parse the remaining part of the dataset and, line by line (thus, while looping over the dataset), execute all the instructions provided.
Instruction parsing & execution
Part 1 requires to move one crate at a time, while part 2 requires to move the specified number of creates at the same time.
Moving one element at a time, means inserting the elements in reversed order (read
CBA). We can keep track of this requirement with a boolean variable.
one_at_a_time = tf.Variable(True)
Now we can focus on the line parsing. As anticipated, we have a very limited number of functions in the
tf.strings package for working with strings, thus the manipulation is a bit cumbersome.
Given a line in the format
move X from A to B
our goal is to extract the numbers
B (destination) and use them for reading from the stack with ID
source, the required
amount of elements from the top (where our “top” is given by the number of nonempty elements in the stack), and move them into
B according to the strategy defined by
def move(line): amount = tf.strings.to_number( tf.strings.regex_replace( tf.strings.regex_replace(line, "move ", ""), r" from \d* to \d*$", "" ), tf.int64, ) source_dest = tf.strings.regex_replace(line, r"move \d* from ", "") source = ( tf.strings.to_number( tf.strings.regex_replace(source_dest, r" to \d*$", ""), tf.int64 ) - 1 ) dest = ( tf.strings.to_number( tf.strings.regex_replace(source_dest, r"\d* to ", ""), tf.int64 ) - 1 ) num_element_source = num_elements.lookup([source]) top = max_stack_size - num_element_source read = stacks[top : top + amount, source] # remove from source indices_x, indices_y = tf.meshgrid(tf.range(top, top + amount), [source]) indices = tf.reshape(tf.stack([indices_x, indices_y], axis=-1), (-1, 2)) updates = tf.reshape(tf.repeat("", amount), (-1, 1)) stacks.assign( tf.tensor_scatter_nd_update(stacks, indices, updates), use_locking=True ) num_element_dest = num_elements.lookup([dest]) top = max_stack_size - num_element_dest - 1 # one a at a time -> reverse if one_at_a_time: insert = tf.reverse(read, axis=) insert = tf.reshape(insert, (-1, 1)) else: insert = tf.reshape(read, (-1, 1)) indices_x, indices_y = tf.meshgrid(tf.range(top - amount + 1, top + 1), [dest]) indices = tf.reshape(tf.stack([indices_x, indices_y], axis=-1), (-1, 2)) stacks.assign( tf.tensor_scatter_nd_update(stacks, indices, insert), use_locking=True ) update_num_elements() return stacks
move function does all the job. Reads the line, parses it, extracts the values, removes them from source (set the to empty string), and inserts into the destination stack the read values according to the strategy.
It’s not immediate to understand all the indices manipulation, so take your time for reading the code carefully. In particular, the reader should focus on the
tf.meshgrid function used to create the various indices along the two dimensions.
Solving the problem
move function works only on the instructions lines, thus we need to create a dataset that only contains these lines. It’s trivial with the
moves_dataset = dataset.skip(tf.shape(stacks_tensor, tf.int64) + 2)
We are now ready to play, solve the problem and visualize the final stack:
tf.print("part 1") play = moves_dataset.map(move) list(play) indices_x = tf.range(num_stacks - 1) indices_y = max_stack_size - tf.reverse(num_elements.export(), axis=) indices = tf.reshape(tf.stack([indices_y, indices_x], axis=-1), (-1, 2)) tf.print(tf.strings.join(tf.squeeze(tf.gather_nd(stacks, indices)), ""))
The thing worth noting is that playing this game is the iteration of the
play dataset, invocated by converting the dataset object to a list (
num_elements lookup table contains the index of every top-crate, thus we can just gather them, join as a single string and print them for getting the expected result.
Part two is identical, we only need to toggle the
one_at_a_time variable, reset the
stacks variable to its initial state, play once again the very same game, and gather the result as above.
tf.print("part 2") initialize_stacks() update_num_elements() one_at_a_time.assign(False) play = moves_dataset.map(move) list(play) indices_x = tf.range(num_stacks - 1) indices_y = max_stack_size - tf.reverse(num_elements.export(), axis=) indices = tf.reshape(tf.stack([indices_y, indices_x], axis=-1), (-1, 2)) tf.print(tf.strings.join(tf.squeeze(tf.gather_nd(stacks, indices)), ""))
Here we go, day 5 problem solved in pure TensorFlow!
tf.Variable undefined shape & validation
In this part of the article, I just want to highlight a not widely used behavior of the
tf.Variable object: the possibility of creating a
tf.Variable without specifying its shape.
tf.Variable documentation mentions the
validate_shape parameter only twice:
- When documenting the
The initial value must have a shape specified unless validate_shape is set to False.
- When documenting the
If False, allows the variable to be initialized with a value of unknown shape. If True, the default, the shape of
initial_valuemust be known.
So, by reading the documentation we can say that’s possible to define a
tf.Variable without shape, assigning it a
tf.Tensor with a different shape without any problem. In this way, we could avoid the limitation of specifying a
max_stack_size and writing a very general solution without this hard constraint.
So, instead of giving a complete solution I just want to leave here a potential starting point for designing a more general solution.
stacks = tf.Variable( stacks_tensor, validate_shape=False, dtype=tf.string, shape=tf.TensorShape(None) ) # ... def initialize_stacks(): stacks.assign(tf.zeros_like(stacks)) shape = tf.shape(stacks) stacks.assign( tf.reshape( stacks, [ shape, shape, 1, ], ) ) # ...
The goal is the remove from the solution the
max_stack_size and always work with a
tf.Variable with a variable shape: the tricky part will be the indexing, I guarantee it!
You can see the complete solution in folder
5 in the dedicated GitHub repository (in the
2022 folder): https://github.com/galeone/tf-aoc.
Using TensorFlow for solving this problem allowed us to understand that the
tf.strings package contains only a small set of utilities for working with strings, and that the regular expression support is quite limited.
Once again, we relied upon the experimental
tf.lookup.experimental.MutableHashTable, although still experimental (by several years!) it works quite well.
In the last part of the article I suggested an alternative approach to the problem, that should allow a better design of this solution without a constraint on the maximum number of stacks (and thus, without a constraint on the variable dimension). If you want to contribute and submit to the repository a merge request with your alternate solution in pure TensorFlow with a
tf.Variable with an undefined shape, I’d be very happy to review and talk about it!
If you missed the article about the previous days’ solutions, here’s a handy list
- Advent of Code 2022 in pure TensorFlow - Days 1 & 2.
- Advent of Code 2022 in pure TensorFlow - Days 3 & 4.
For any feedback or comment, please use the Disqus form below - thanks!