def __init__(
      self,
      remote_device,
      local_device,
      top_delta_size=64,
      top_delta_layers=2,
      compute_h_size=64,
      compute_h_layers=1,
      delta_dim=32,
      num_grad_channels=4,
      normalize_epsilon=1.,
  ):
    self.local_device = local_device
    self.remote_device = remote_device
    self.top_delta_size = top_delta_size
    self.top_delta_layers = top_delta_layers
    self.compute_h_size = compute_h_size
    self.compute_h_layers = compute_h_layers
    self.delta_dim = delta_dim
    self.num_grad_channels = num_grad_channels
    self.normalize_epsilon = normalize_epsilon,

    with tf.device(local_device):
      self.opt = optimizers.UnrollableGradientDescentRollingOptimizer(
          learning_rate=1e-4)

    # lazily initialized for readouts
    self.readout_mods = {}

    super(MoreLocalWeightUpdateProcess,
          self).__init__(name='MoreLocalWeightUpdateProcess')

    with tf.device(remote_device):
      self()

 def _apply_drop_path(self, net):
   """Apply drop_path regularization to net."""
   drop_path_keep_prob = self._drop_path_keep_prob
   if drop_path_keep_prob < 1.0:
     # Scale keep prob by layer number
     assert self._cell_num != -1
     # The added 2 is for the reduction cells
     num_cells = self._total_num_cells
     layer_ratio = (self._cell_num + 1)/float(num_cells)
     with tf.device('/cpu:0'):
       tf.summary.scalar('layer_ratio', layer_ratio)
     drop_path_keep_prob = 1 - layer_ratio * (1 - drop_path_keep_prob)
     # Decrease the keep probability over time
     current_step = tf.cast(tf.train.get_or_create_global_step(),
                            tf.float32)
     drop_path_burn_in_steps = self._total_training_steps
     current_ratio = (
         current_step / drop_path_burn_in_steps)
     current_ratio = tf.minimum(1.0, current_ratio)
     with tf.device('/cpu:0'):
       tf.summary.scalar('current_ratio', current_ratio)
     drop_path_keep_prob = (
         1 - current_ratio * (1 - drop_path_keep_prob))
     with tf.device('/cpu:0'):
       tf.summary.scalar('drop_path_keep_prob', drop_path_keep_prob)
     net = drop_path(net, drop_path_keep_prob)
   return net

    def __init__(self, session, np_matrix, rank,
                 learning_rate=0.1):
        matrix = tf.constant(np_matrix, dtype=tf.float32)
        scale = 2 * np.sqrt(np_matrix.mean() / rank)
        initializer = tf.random_uniform_initializer(maxval=scale)

        with tf.device('/job:ps/task:0'):
            self.matrix_W = tf.get_variable(
                "W", (np_matrix.shape[0], rank), initializer=initializer
            )
        with tf.device("/job:ps/task:1"):
            self.matrix_H = tf.get_variable(
                "H", (rank, np_matrix.shape[1]), initializer=initializer
            )

        matrix_WH = tf.matmul(self.matrix_W, self.matrix_H)
        f_norm = tf.reduce_sum(tf.pow(matrix - matrix_WH, 2))

        nn_w = tf.reduce_sum(tf.abs(self.matrix_W) - self.matrix_W)
        nn_h = tf.reduce_sum(tf.abs(self.matrix_H) - self.matrix_H)
        constraint = INFINITY * (nn_w + nn_h)
        self.loss = f_norm + constraint
        self.constraint = constraint

        self.session = session
        self.optimizer = tf.train.GradientDescentOptimizer(
            learning_rate
        ).minimize(self.loss)

def prepare_networks(gpu,image_batch, nb_cl, nb_groups):
  mean_img = tf.constant([123.68, 116.779, 103.939], dtype=tf.float32, shape=[1, 1, 1, 3], name='img_mean')
  scores   = []
  with tf.variable_scope('ResNet18'):
    with tf.device('/gpu:' + gpu):
        score = utils_resnet.ResNet18(image_batch-mean_img, phase='train',num_outputs=nb_cl*nb_groups)
        scores.append(score)
    
    scope = tf.get_variable_scope()
    scope.reuse_variables()
  
  # First score and initialization
  variables_graph = tf.get_collection(tf.GraphKeys.WEIGHTS, scope='ResNet18')
  scores_stored   = []
  with tf.variable_scope('store_ResNet18'):
    with tf.device('/gpu:' + gpu):
        score = utils_resnet.ResNet18(image_batch-mean_img, phase='test',num_outputs=nb_cl*nb_groups)
        scores_stored.append(score)
    
    scope = tf.get_variable_scope()
    scope.reuse_variables()
  
  variables_graph2 = tf.get_collection(tf.GraphKeys.WEIGHTS, scope='store_ResNet18')
  
  return variables_graph,variables_graph2,scores,scores_stored

  def testAnalysisAndAllocations(self):
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    config = tf.ConfigProto(device_count={'CPU': 3})

    with tf.Session(config=config) as sess:
      with tf.device('/cpu:0'):
        const1 = tf.constant(1.0, name='const1')
      with tf.device('/cpu:1'):
        const2 = tf.constant(2.0, name='const2')
      with tf.device('/cpu:2'):
        result = const1 + const2 + const1 * const2
      sess.run(result, options=run_options, run_metadata=run_metadata)

    self.assertTrue(run_metadata.HasField('step_stats'))
    tl = timeline.Timeline(run_metadata.step_stats)
    step_analysis = tl.analyze_step_stats()
    ctf = step_analysis.chrome_trace.format_to_string()
    self._validateTrace(ctf)
    maximums = step_analysis.allocator_maximums
    self.assertTrue('cpu' in maximums)
    cpu_max = maximums['cpu']
    # At least const1 + const2, both float32s (4 bytes each)
    self.assertGreater(cpu_max.num_bytes, 8)
    self.assertGreater(cpu_max.timestamp, 0)
    self.assertTrue('const1' in cpu_max.tensors)
    self.assertTrue('const2' in cpu_max.tensors)

  def testReturnsSingleCheckpointIfOneShardedCheckpoint(self):
    checkpoint_dir = os.path.join(self.get_temp_dir(),
                                  'one_checkpoint_found_sharded')
    if not tf.gfile.Exists(checkpoint_dir):
      tf.gfile.MakeDirs(checkpoint_dir)

    global_step = tf.contrib.framework.get_or_create_global_step()

    # This will result in 3 different checkpoint shard files.
    with tf.device('/cpu:0'):
      tf.Variable(10, name='v0')
    with tf.device('/cpu:1'):
      tf.Variable(20, name='v1')

    saver = tf.train.Saver(sharded=True)

    with tf.Session(
        target='',
        config=tf.ConfigProto(device_count={'CPU': 2})) as session:

      session.run(tf.initialize_all_variables())
      save_path = os.path.join(checkpoint_dir, 'model.ckpt')
      saver.save(session, save_path, global_step=global_step)

    num_found = 0
    for _ in tf.contrib.training.checkpoints_iterator(
        checkpoint_dir, timeout=0):
      num_found += 1
    self.assertEqual(num_found, 1)

def pack_range(key, packing, grad_vars, rng):
    """Form the concatenation of a specified range of gradient tensors.

  Args:
    key: Value under which to store meta-data in packing that will be used
      later to restore the grad_var list structure.
    packing: Dict holding data describing packed ranges of small tensors.
    grad_vars: List of (grad, var) pairs for one tower.
    rng: A pair of integers giving the first, last indices of a consecutive
      range of tensors to be packed.

  Returns:
    A tensor that is the concatenation of all the specified small tensors.
  """
    to_pack = grad_vars[rng[0]:rng[1] + 1]
    members = []
    variables = []
    restore_shapes = []
    with tf.name_scope('pack'):
        for g, v in to_pack:
            variables.append(v)
            restore_shapes.append(g.shape)
            with tf.device(g.device):
                members.append(tf.reshape(g, [-1]))
        packing[key] = GradPackTuple(
            indices=range(rng[0], rng[1] + 1),
            vars=variables,
            shapes=restore_shapes)
        with tf.device(members[0].device):
            return tf.concat(members, 0)

def extract_features(ids, path, output_path, extractor, batch_size=64):
    images_names = dict()
    for p in listdir(path):
        image_id = int(p.split('_')[-1].split('.')[0])
        if image_id in ids:
            images_names[image_id] = p
    batch,names = [],[]
    with open(output_path,'w') as output_file:
        for idx,n in enumerate(images_names):
            p = join(path, images_names[n])
            batch.append(load_image(p))
            names.append(n)
            if len(batch)==batch_size:
                batch = np.stack(batch)
                feed_dict = {images: batch}
                with tf.device('/gpu:0'):
                    features = sess.run(extractor, feed_dict=feed_dict)
                for n,f in zip(names,features):
                    output_file.write("%s;%s\n" % (n, " ".join(str(x) for x in f)))
                print("%d/%d" % (idx,len(images_names)))
                batch, names = [],[]
                output_file.flush()
        if len(batch)>0:
            batch = np.stack(batch)
            feed_dict = {images: batch}
            with tf.device('/gpu:0'):
                features = sess.run(extractor, feed_dict=feed_dict)
            for n,f in zip(names,features):
                output_file.write("%s;%s\n" % (n, " ".join(str(x) for x in f)))
            print("%d/%d" % (idx,len(images_names)))
            output_file.flush()

  def __call__(self, getter, name, *args, **kwargs):
    staging_ops = self.variable_mgr.staging_vars_on_devices[self.device_num]
    if name in staging_ops:
      put_op, get_op = staging_ops[name]
      return get_op
    real_var = getter(name, *args, **kwargs)
    shape = kwargs['shape']
    dtype = kwargs['dtype']
    trainable = kwargs['trainable']
    if self.cpu_device:
      with tf.device(self.cpu_device):
        # This helps copying the weights from the parameter to this server only
        # once.
        if name in self.variable_mgr.staged_vars_on_cpu:
          cpu_var = self.variable_mgr.staged_vars_on_cpu[name]
        else:
          cpu_var = tf.identity(real_var)
          self.variable_mgr.staged_vars_on_cpu[name] = cpu_var
      var_to_stage = cpu_var
    else:
      var_to_stage = tf.identity(real_var)  # de-reference the variable.

    with tf.device(self.devices[self.device_num]):
      staging_area = data_flow_ops.StagingArea([dtype], shapes=[shape])
      put_op = staging_area.put([var_to_stage])
      get_op = staging_area.get()[0]
      staging_ops[name] = (put_op, get_op)
    if trainable:
      # For trainable variables, they are managed separatedly through
      # apply_gradients.
      return get_op
    else:
      # For other shadow variables, the access is decoupled through a wrapper
      # class.
      return StagedModelVariable(real_var, get_op, self.variable_mgr)

def create_weight_variables(shape, seed, name, use_gpu=False):
    """
    Create gaussian random neurons with mean 0 and std 0.1

    **Paramters**

      shape: Shape of the layer
    """

    #import ipdb; ipdb.set_trace()

    if len(shape) == 4:
        in_out = shape[0] * shape[1] * shape[2] + shape[3]
    else:
        in_out = shape[0] + shape[1]

    import math
    stddev = math.sqrt(3.0 / in_out) # XAVIER INITIALIZER (GAUSSIAN)

    initializer = tf.truncated_normal(shape, stddev=stddev, seed=seed)
    
    if use_gpu:
        with tf.device("/gpu"):
            return tf.get_variable(name, initializer=initializer, dtype=tf.float32)
    else:
        with tf.device("/cpu"):
            return tf.get_variable(name, initializer=initializer, dtype=tf.float32)

def all_sync_params(tower_params, devices, usenccl=True):
    """Assigns the params from the first tower to all others"""
    if len(devices) == 1:
        return tf.no_op()
    sync_ops = []
    if have_nccl and usenccl:
        for param_on_devices in zip(*tower_params):
            # print('PARAM_ON_DEVICES: {}'.format(param_on_devices))  # DEBUG
            # Note: param_on_devices is [paramX_gpu0, paramX_gpu1, ...]
            param0 = param_on_devices[0]
            send_op, received_tensors = nccl.broadcast(param0, devices[1:])
            sync_ops.append(send_op)
            for device, param, received in zip(devices[1:],
                                               param_on_devices[1:],
                                               received_tensors):
                with tf.device(device):
                    sync_op = param.assign(received)
                    sync_ops.append(sync_op)
    else:
        params0 = tower_params[0]
        for device, params in zip(devices, tower_params):
            with tf.device(device):
                for param, param0 in zip(params, params0):
                    sync_op = param.assign(param0.read_value())
                    sync_ops.append(sync_op)

    return tf.group(*sync_ops)

def train():
    hyperparams = {'batch_size': 50,
                   'learning_rate': 0.0001,
                   'grad_decay': 0.95,
                   'grad_epsilon': 0.01,
                   'num_updates': 20000,
                   'grad_norm_clip': 5}
    with tf.device('/cpu:0'):
        model = TradingSystemsModel(hyperparams)
    loss = tb.Crossentropy(hyperparams)
    acc = tb.CatAcc(hyperparams)
    evaluator = tb.Evaluator(hyperparams, loss, acc)
    optim = tb.RMSPropOptim(hyperparams)
    trainer = tb.Trainer(model, hyperparams, loss, optim, evaluator)

    split = 90000
    data = np.load('data/trading-systems.npz')
    print(data['ticks'].shape)
    train_xs = {'ticks': data['ticks'][:split]}
    train_y = data['targets'][:split]

    val_xs = {'ticks': data['ticks'][split:]}
    val_y = data['targets'][split:]

    with tf.device('/cpu:0'):
        trainer.train(train_xs, train_y,
                      val_xs, val_y,
                      val_cmp=True)
    evaluator.eval(model, val_xs, val_y)

def all_reduce_gradients(tower_grads, devices):
    average_grads = []
    for grad_and_vars in zip(*tower_grads):
        # Note that each grad_and_vars looks like the following:
        #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
        grads = []
        split_grads = []
        assert len(grad_and_vars) == FLAGS.num_workers
        # Each GPU splits its own grad
        for i, (g, _) in enumerate(grad_and_vars):
            with tf.device(devices[i]):
                split_grads.append(tf.split(0, FLAGS.num_workers, g))
        # Each GPU gatheres slices of grad from other GPUs to do average.
        for i, dev in enumerate(devices):
            with tf.device(dev):
                x = split_grads[i][i]
                for j in range(FLAGS.num_workers):
                    if i == j:
                        continue
                    x += split_grads[j][i]
                grads.append(x / FLAGS.num_workers)
        grad = tf.concat(0, grads)

        # Keep in mind that the Variables are redundant because they are shared
        # across towers. So .. we will just return the first tower's pointer to
        # the Variable.
        v = grad_and_vars[0][1]
        grad_and_var = (grad, v)
        average_grads.append(grad_and_var)
    return average_grads

  def _add_shared_train_op(self):
    """Sets self._train_op, the op to run for training."""
    # Take gradients of the trainable variables w.r.t. the loss function to minimize
    if self._hps.rl_training or self._hps.ac_training:
      loss_to_minimize = self._reinforce_shared_loss
      if self._hps.coverage:
        loss_to_minimize = self._reinforce_cov_total_loss
    else:
      loss_to_minimize = self._pgen_loss
      if self._hps.coverage:
        loss_to_minimize = self._pointer_cov_total_loss

    tvars = tf.trainable_variables()
    gradients = tf.gradients(loss_to_minimize, tvars, aggregation_method=tf.AggregationMethod.EXPERIMENTAL_TREE)

    # Clip the gradients
    with tf.device("/gpu:{}".format(self._hps.gpu_num)):
      grads, global_norm = tf.clip_by_global_norm(gradients, self._hps.max_grad_norm)

    # Add a summary
    tf.summary.scalar('global_norm', global_norm)

    # Apply adagrad optimizer
    optimizer = tf.train.AdagradOptimizer(self._hps.lr, initial_accumulator_value=self._hps.adagrad_init_acc)
    with tf.device("/gpu:{}".format(self._hps.gpu_num)):
      self._shared_train_op = optimizer.apply_gradients(zip(grads, tvars), global_step=self.global_step, name='train_step')

    def _build_word_embeddings(self):
        n_tokens_vocab = self.options['n_tokens_vocab']
        batch_size = self.options['batch_size']
        unroll_steps = self.options['unroll_steps']

        # LSTM options
        projection_dim = self.options['lstm']['projection_dim']

        # the input token_ids and word embeddings
        self.token_ids = tf.placeholder(DTYPE_INT,
                               shape=(batch_size, unroll_steps),
                               name='token_ids')
        # the word embeddings
        with tf.device("/cpu:0"):
            self.embedding_weights = tf.get_variable(
                "embedding", [n_tokens_vocab, projection_dim],
                dtype=DTYPE,
            )
            self.embedding = tf.nn.embedding_lookup(self.embedding_weights,
                                                self.token_ids)

        # if a bidirectional LM then make placeholders for reverse
        # model and embeddings
        if self.bidirectional:
            self.token_ids_reverse = tf.placeholder(DTYPE_INT,
                               shape=(batch_size, unroll_steps),
                               name='token_ids_reverse')
            with tf.device("/cpu:0"):
                self.embedding_reverse = tf.nn.embedding_lookup(
                    self.embedding_weights, self.token_ids_reverse)

def main(unused_argv):
  tf.logging.set_verbosity(FLAGS.log)

  if not tf.gfile.Exists(FLAGS.logdir):
    tf.gfile.MakeDirs(FLAGS.logdir)

  with tf.Graph().as_default():

    # If ps_tasks is 0, the local device is used. When using multiple
    # (non-local) replicas, the ReplicaDeviceSetter distributes the variables
    # across the different devices.
    model = utils.get_module("baseline.models.%s" % FLAGS.model)
    hparams = model.get_hparams(FLAGS.config)

    # Run the Reader on the CPU
    cpu_device = ("/job:worker/cpu:0" if FLAGS.ps_tasks else
                  "/job:localhost/replica:0/task:0/cpu:0")

    with tf.device(cpu_device):
      with tf.name_scope("Reader"):
        batch = reader.NSynthDataset(
            FLAGS.train_path, is_training=True).get_baseline_batch(hparams)

    with tf.device(tf.train.replica_device_setter(ps_tasks=FLAGS.ps_tasks)):
      train_op = model.train_op(batch, hparams, FLAGS.config)

      # Run training
      slim.learning.train(
          train_op=train_op,
          logdir=FLAGS.logdir,
          master=FLAGS.master,
          is_chief=FLAGS.task == 0,
          number_of_steps=hparams.max_steps,
          save_summaries_secs=FLAGS.save_summaries_secs,
          save_interval_secs=FLAGS.save_interval_secs)

 def testManyCPUs(self):
   run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
   run_metadata = tf.RunMetadata()
   config = tf.ConfigProto(device_count={'CPU': 3})
   with tf.Session(config=config) as sess:
     with tf.device('/cpu:0'):
       const1 = tf.constant(1.0, name='const1')
     with tf.device('/cpu:1'):
       const2 = tf.constant(2.0, name='const2')
     with tf.device('/cpu:2'):
       result = const1 + const2 + const1 * const2
     sess.run(result, options=run_options, run_metadata=run_metadata)
   self.assertTrue(run_metadata.HasField('step_stats'))
   step_stats = run_metadata.step_stats
   devices = [d.device for d in step_stats.dev_stats]
   self.assertTrue('/job:localhost/replica:0/task:0/cpu:0' in devices)
   self.assertTrue('/job:localhost/replica:0/task:0/cpu:1' in devices)
   self.assertTrue('/job:localhost/replica:0/task:0/cpu:2' in devices)
   tl = timeline.Timeline(step_stats)
   ctf = tl.generate_chrome_trace_format()
   self._validateTrace(ctf)
   tl = timeline.Timeline(step_stats)
   ctf = tl.generate_chrome_trace_format(show_dataflow=False)
   self._validateTrace(ctf)
   tl = timeline.Timeline(step_stats)
   ctf = tl.generate_chrome_trace_format(show_memory=False)
   self._validateTrace(ctf)
   tl = timeline.Timeline(step_stats)
   ctf = tl.generate_chrome_trace_format(show_memory=False,
                                         show_dataflow=False)
   self._validateTrace(ctf)

    def build_generator(self):
	
	# placeholder is for feeding data
	image = tf.placeholder(tf.float32, [self.batch_size, self.dim_image])  # (batch_size, dim_image)
	local_image = tf.placeholder(tf.float32, [self.batch_size, self.dim_image])
	query = tf.placeholder(tf.int32, [self.batch_size, MAX_QUERY_WORDS])
	query_mask = tf.placeholder(tf.float32, [self.batch_size, MAX_QUERY_WORDS])
	bbox = tf.placeholder(tf.float32, [self.batch_size, self.dim_coordinates])

	# [image] embed image feature to dim_hidden
        image_emb = tf.nn.bias_add(tf.matmul(image, self.embed_image_W), self.embed_image_b) # (batch_size, dim_hidden)
	local_image_emb = tf.nn.bias_add(tf.matmul(local_image, self.embed_local_W), self.embed_local_b) # (batch_size, dim_hidden)
	
        score = tf.zeros([self.batch_size], tf.float32)

	state_lang = tf.zeros([self.batch_size, self.lstm_lang.state_size])
	state_context = tf.zeros([self.batch_size, self.lstm_context.state_size])
	state_local = tf.zeros([self.batch_size, self.lstm_local.state_size])
	query_emb = tf.zeros([self.batch_size, self.dim_hidden])
	for j in range(MAX_QUERY_WORDS): 


	    # language lstm
            with tf.variable_scope("lstm_lang"):
                output_lang, state_lang = self.lstm_lang(query_emb, state_lang)
            lang = tf.slice(state_lang, [0,0], [self.batch_size, self.dim_hidden])
            # context lstm

            with tf.variable_scope("lstm_context"):
                output_context, state_context = self.lstm_context(tf.concat(1,[image_emb, lang]), state_context)
            context = tf.slice(state_context, [0,0], [self.batch_size, self.dim_hidden])

            # local lstm
            with tf.variable_scope("lstm_local"):
                output_local, state_local = self.lstm_local(tf.concat(1,[local_image_emb, lang, bbox]), state_local)
            local = tf.slice(state_local, [0,0], [self.batch_size, self.dim_hidden])

            context_emb = tf.nn.xw_plus_b(context, self.W_context, self.B_context)
            local_emb = tf.nn.xw_plus_b(local, self.W_local, self.B_local)
            word_pred = tf.add(context_emb, local_emb)

	    max_prob_index = tf.argmax(word_pred, 1) # b

	    labels = tf.expand_dims(query[:,j], 1)
            indices = tf.expand_dims(tf.range(0, self.batch_size, 1), 1)
            concated = tf.concat(1, [indices, labels])
            with tf.device('/cpu:0'):
                onehot_labels = tf.sparse_to_dense(concated, tf.pack([self.batch_size, self.dict_words]), 1.0, 0.0)
	    current_score = tf.mul(onehot_labels, word_pred)
	    current_score = tf.reduce_sum(current_score, 1)
	    current_score = tf.mul(current_score, query_mask[:,j])
	    current_score = tf.reshape(current_score, [1,self.batch_size])
	    current_score = tf.nn.softmax(current_score)
	    score = tf.add(score, current_score)

            with tf.device("/cpu:0"):
                tf.get_variable_scope().reuse_variables()
                query_emb = tf.nn.embedding_lookup(self.query_emb_W, max_prob_index)

	return score, image, local_image, query, query_mask, bbox

    def test_single_output(self):
        print('*** Running Test: ' + self.__class__.__name__ + ' function: ' + _getframe().f_code.co_name)

        class AddOp(Operator):
            def op(self, x, y):
                pos = position_in(x.shape)
                out = output_like(x)
                out[pos] = x[pos] + y[pos]
                return out

        in0 = np.random.random(5).astype(np.float32)
        in1 = np.random.random(5).astype(np.float32)
        reference = 4*(in0 + in1)*(in0 + in1)

        with tf.Session() as sess:
            with tf.device('/cpu:0'):
                a = in0*2
                b = in1*2
                c = AddOp(a, b, clear_cache=True).as_tensorflow()
                squared = tf.square(c)
            if cuda_enabled:
                with tf.device('/gpu:0'):
                    a_gpu = in0*2
                    b_gpu = in1*2
                    c_gpu = AddOp(a_gpu, b_gpu).as_tensorflow()
                    squared_gpu = tf.square(c_gpu)
                result, result_gpu = sess.run([squared, squared_gpu])
                assert np.allclose(reference, result_gpu)
            else:
                result = sess.run([squared])

        assert np.allclose(reference, result)

    def test_single_output(self):

        @operator()
        def add(x, y):
            pos = position_in(x.shape)
            out = output_like(x)
            out[pos] = x[pos] + y[pos]
            return out

        in0 = np.random.random(5).astype(np.float32)
        in1 = np.random.random(5).astype(np.float32)
        reference = 4*(in0 + in1)*(in0 + in1)

        test_config = tf.ConfigProto(allow_soft_placement=False)
        # Don't perform optimizations for tests so we don't inadvertently run
        # gpu ops on cpu
        test_config.graph_options.optimizer_options.opt_level = -1
        with tf.Session(config=test_config) as sess:
            with tf.device('/cpu:0'):
                a = in0*2
                b = in1*2
                c = as_tensorflow(add(a, b))
                squared = tf.square(c)
            if cuda_enabled:
                with tf.device('/gpu:0'):
                    a_gpu = in0*2
                    b_gpu = in1*2
                    c_gpu = as_tensorflow(add(a_gpu, b_gpu))
                    squared_gpu = tf.square(c_gpu)
                result, result_gpu = sess.run([squared, squared_gpu])
                assert np.allclose(reference, result_gpu)
            else:
                result = sess.run([squared])

        assert np.allclose(reference, result)