Collatz with AutoGraph

collatz.py

def collatz(a):
    counter = 0
    while a != 1:
        if a % 2 == 0:
            a = a // 2
        else:
            a = 3 * a + 1
        counter = counter + 1
    return counter

graph_mode_collatz = autograph.to_graph(collatz)
# The code is human-readable, too
print(autograph.to_code(collatz))

collatz_tensor = graph_mode_collatz(tf.constant(n))