# Path: morph_net/framework/tpu_util.py
# def run_on_cpu():
# def is_on_cpu():
# def get_variable_name(read_variable_op):
# def maybe_convert_to_variable(tensor):
# def write_to_variable(tensor, fail_if_exists=True):
# def read_from_variable(tensor):
# def is_intermediate_var(v):


from absl.testing import parameterized
from morph_net.framework import tpu_util
from tensorflow.contrib import slim
import tensorflow.compat.v1 as tf

  def test_write_to_variable(self):
    foo = tf.constant(0., name='foo')
    tpu_util.write_to_variable(foo)

    with self.assertRaises(ValueError):
      tpu_util.write_to_variable(foo, fail_if_exists=True)

    # Variable sharing behavior should be dictated by `fail_if_exists` which
    # overrides the effect of outer scopes.
    with tf.variable_scope('', reuse=True):
      # should fail to return existing variable even though reuse=True
      with self.assertRaises(ValueError):
        tpu_util.write_to_variable(foo, fail_if_exists=True)

    with tf.variable_scope('', reuse=False):
      # should return existing variable even though reuse=False
      foo_copy = tpu_util.write_to_variable(foo, fail_if_exists=False)
      self.assertEqual(tpu_util.var_store[foo], tpu_util.var_store[foo_copy])
      self.assertLen(set(tpu_util.var_store.values()), 1)

    with tf.variable_scope('', reuse=True):
      # should create new variable even though reuse=True
      bar = tf.constant(0., name='bar')
      tpu_util.write_to_variable(bar)
      self.assertLen(set(tpu_util.var_store.values()), 2)

  def test_write_to_variable_check_inputs(self):
    variable = tf.get_variable('x', shape=[1], dtype=tf.float32)
    with self.assertRaises(ValueError):

# Path: morph_net/framework/tpu_util.py
# def run_on_cpu():
# def is_on_cpu():
# def get_variable_name(read_variable_op):
# def maybe_convert_to_variable(tensor):
# def write_to_variable(tensor, fail_if_exists=True):
# def read_from_variable(tensor):
# def is_intermediate_var(v):


# Path: morph_net/framework/tpu_util_test.py
from absl.testing import parameterized
from morph_net.framework import tpu_util
from tensorflow.contrib import slim
import tensorflow.compat.v1 as tf


  def test_write_to_variable(self):
    foo = tf.constant(0., name='foo')
    tpu_util.write_to_variable(foo)

    with self.assertRaises(ValueError):
      tpu_util.write_to_variable(foo, fail_if_exists=True)

    # Variable sharing behavior should be dictated by `fail_if_exists` which
    # overrides the effect of outer scopes.
    with tf.variable_scope('', reuse=True):
      # should fail to return existing variable even though reuse=True
      with self.assertRaises(ValueError):
        tpu_util.write_to_variable(foo, fail_if_exists=True)

    with tf.variable_scope('', reuse=False):
      # should return existing variable even though reuse=False
      foo_copy = tpu_util.write_to_variable(foo, fail_if_exists=False)
      self.assertEqual(tpu_util.var_store[foo], tpu_util.var_store[foo_copy])
      self.assertLen(set(tpu_util.var_store.values()), 1)

    with tf.variable_scope('', reuse=True):
      # should create new variable even though reuse=True
      bar = tf.constant(0., name='bar')
      tpu_util.write_to_variable(bar)
      self.assertLen(set(tpu_util.var_store.values()), 2)

  def test_write_to_variable_check_inputs(self):
    variable = tf.get_variable('x', shape=[1], dtype=tf.float32)
    with self.assertRaises(ValueError):

      tpu_util.write_to_variable(variable)