loop-diff-cpu.py

import tensorflow as tf
from tensorflow.contrib.compiler import xla
import numpy as np


a = tf.Variable(1.0, use_resource=True)

def repeat(count, body, vars):
  return tf.while_loop(lambda *args : True, body, vars, maximum_iterations=count)

def model(x):
  def body(x):
    return a*x
  return repeat(5, body, [x])


def run():
  def build(features):
    logits = model(features)
    loss = tf.math.reduce_sum(logits)
    optimizer = tf.train.MomentumOptimizer(learning_rate=.001, momentum=0.9)
    grads_and_vars = optimizer.compute_gradients(loss)
    train_op = optimizer.apply_gradients(grads_and_vars)
    return loss, train_op

  with tf.device('cpu'):
    features = tf.placeholder(tf.float32, shape=[10])
    compiled = xla.compile(build, inputs=[features])

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(compiled, feed_dict = { features: np.random.random(10) }))


if __name__ == '__main__':
  run()