Script to create graph of toy matrix multiplication and add dependencies after creating the graph

toy_matmul_graph.py

import tensorflow as tf
import json
import sys
from tensorflow.python.client import timeline
from protobuf_to_dict import protobuf_to_dict
from google.protobuf.json_format import MessageToJson
from tensorflow.core.protobuf import rewriter_config_pb2

dim = 16384

with tf.device('/gpu:0'):
    X9 = tf.random_uniform([dim, dim], 0, 10)
    _X9 = tf.placeholder(dtype=tf.float32, shape=[dim, dim])
    Z19 = tf.matmul(_X9, _X9)

with tf.device('/gpu:1'):
    Y9 = tf.random_uniform([dim, dim], 0, 10)
    _Y9 = tf.placeholder(dtype=tf.float32, shape=[dim, dim])
    Z29 = tf.matmul(_Y9, _Y9)

    W9 = tf.matmul(Z19, Z29)

n = 8
dim = 32
Z1, Z2, W = [], [], []
X, _X, Y, _Y = [], [], [], []

with tf.device('/gpu:0'):
    for i in range(n):
        dim *= 2

        X.append(tf.random_uniform([dim, dim], 0, 10, name='X' + str(i)))
        Y.append(tf.random_uniform([dim, dim], 0, 10, name='Y' + str(i)))

        # with tf.control_dependencies([Z19]):
        _X.append(tf.placeholder(dtype=tf.float32, shape=[dim, dim]))
        _X[i].op.control_inputs.append(Z19.op)
        Z1.append(tf.matmul(_X[i], _X[i]))
        Z1[i].op.control_inputs.append(Z19.op)
        # print(Z1[i].op.control_inputs)

        _Y.append(tf.placeholder(dtype=tf.float32, shape=[dim, dim]))
        _Y[i].op.control_inputs.append(Z19.op)
        Z2.append(tf.matmul(_Y[i], _Y[i]))
        Z2[i].op.control_inputs.append(Z19.op)
        
        W.append(tf.matmul(Z1[i], Z2[i]))
        print(Z1[i].op.control_inputs)
        print(Z2[i].op.control_inputs)


tf.train.export_meta_graph('toy_matmul_graph.meta')