def testMultipleApplicationsInDifferentScopes(self):
    export_path = os.path.join(self.get_temp_dir(), "module-applied-in-scope")

    spec = hub.create_module_spec(another_stateful_module_fn)
    stateful_module = hub.Module(spec, name="moduleA")
    with tf.name_scope("foo"):
      with tf.variable_scope("bar"):
        times2 = stateful_module(tf.constant([2.0]))
    with tf.name_scope("baz"):
      times3 = stateful_module(tf.constant([3.0]))

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(sess.run(times2), [6.0])
      self.assertAllClose(sess.run(times3), [9.0])
      self.assertEqual(len(stateful_module.variable_map), 1)
      self.assertEquals(
          stateful_module.variable_map["iamtheoneandonly"].name,
          "moduleA/iamtheoneandonly:0")
      stateful_module.export(export_path, sess)

    # Check minimal functionality of the exported module.
    tf.reset_default_graph()
    stateful_module = hub.Module(export_path, name="moduleB")
    times2 = stateful_module(tf.constant([2.0]))
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(sess.run(times2), [6.0])

  def testApplyStatefulModuleMultipleTimes(self):
    export_path = os.path.join(self.get_temp_dir(), "another-module")

    with tf.Session() as sess:
      spec = hub.create_module_spec(another_stateful_module_fn)
      stateful_module = hub.Module(spec, trainable=True)
      times2 = stateful_module(tf.constant([2.0]))
      times3 = stateful_module(tf.constant([3.0]))
      step = tf.Variable(0, trainable=False, name="global_step")
      # Training will adapt the hidden variable to be approximately 2:
      train = tf.contrib.layers.optimize_loss(
          loss=tf.losses.mean_squared_error(times2, [4.0]),
          global_step=step,
          learning_rate=0.05,
          optimizer="SGD")
      sess.run(tf.global_variables_initializer())
      for _ in range(50):
        sess.run(train)
      self.assertAllClose(sess.run(times2), [4.0])
      self.assertAllClose(sess.run(times3), [6.0])
      stateful_module.export(export_path, sess)
    with tf.Session() as sess:
      stateful_module = hub.Module(export_path)
      times4 = stateful_module(tf.constant([4.0]))
      times5 = stateful_module(tf.constant([5.0]))
      sess.run(tf.global_variables_initializer())
      self.assertAllClose(sess.run(times4), [8.0])
      self.assertAllClose(sess.run(times5), [10.0])

  def testClearControlDependenciesForModuleStateButNotApplyGraphs(self):
    module_spec = hub.create_module_spec(stateless_module_fn)

    with tf.Graph().as_default() as g1:
      v = tf.placeholder(dtype=tf.int64, name="v")
      m = hub.Module(module_spec)
      m(v)

    with tf.Graph().as_default() as g2:
      v = tf.placeholder(dtype=tf.int64, name="v")
      with tf.control_dependencies([v]):
        m = hub.Module(module_spec)
      m(v)

    self.assertEqual(g1.as_graph_def(), g2.as_graph_def())

    with tf.Graph().as_default() as g3:
      v = tf.placeholder(dtype=tf.int64, name="v")
      m = hub.Module(module_spec)
      m(v)

    with tf.Graph().as_default() as g4:
      v = tf.placeholder(dtype=tf.int64, name="v")
      m = hub.Module(module_spec)
      with tf.control_dependencies([v]):
        m(v)

    self.assertNotEqual(g3.as_graph_def(), g4.as_graph_def())

  def testModuleWithRegularizedLayers(self):
    # The linear map y = Mx + b with L2 regularization on M and b
    # when trained at x = [1,1] with L2 loss towards the target y' = [4,4]
    # learns M = [[1,1],[1,1]], b = [1,1], y = [3,3], with eight balanced
    # loss terms: the elements of M, b, and y' - y are all distance 1 from zero.
    train_input = [[1.0, 1.0]]
    target = [[4.0, 4.0]]

    spec = hub.create_module_spec(layers_module_fn)
    with tf.Graph().as_default():
      m = hub.Module(spec, trainable=True)
      x = tf.placeholder(dtype=tf.float32)
      y = m(x)
      squared_loss = tf.losses.mean_squared_error(y, target, weights=2.0)
      # Recover REGULARIZATION_LOSSES from the module.
      total_loss = squared_loss + tf.losses.get_regularization_loss()
      step = tf.Variable(0, trainable=False, name="global_step")
      train = tf.contrib.layers.optimize_loss(
          loss=total_loss, global_step=step, learning_rate=0.1, optimizer="SGD")
      with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for _ in range(50):
          sess.run(train, feed_dict={x: train_input})
        # Verify M = [[1,1],[1,1]], b = [1,1] by evaluating at three points.
        # Without regularization, the result would be an underdetermined mess.
        out = sess.run(y, feed_dict={x: [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]})
        self.assertAllClose(
            out, [[1.0, 1.0], [2.0, 2.0], [2.0, 2.0]], atol=0.001)

  def testUpdateOps(self):
    spec = hub.create_module_spec(update_ops_module_fn)
    with tf.Session() as sess:
      trainable_module = hub.Module(spec, trainable=True)
      fixed_module = hub.Module(spec, trainable=False)

      # TODO(b/62433105): Understand what is the desired behaviour of UPDATE_OPS
      # and applying a Module multiple times. For now UPDATE_OPS probably only
      # do something reasonable if each Module is applied exactly one time.
      trainable_module()
      fixed_module()

      variable = tf.Variable(0.0)
      step = tf.Variable(0, trainable=False, name="global_step")
      train_op = tf.contrib.layers.optimize_loss(
          variable,
          global_step=step,
          learning_rate=0.1,
          optimizer="SGD")

      sess.run(tf.global_variables_initializer())
      sess.run(train_op)
      trainable_module_vars = list(trainable_module.variable_map.values())
      self.assertEqual(len(trainable_module_vars), 1)
      self.assertEqual(sess.run(trainable_module_vars[0]), 1)
      fixed_module_vars = list(fixed_module.variable_map.values())
      self.assertEqual(len(fixed_module_vars), 1)
      self.assertEqual(sess.run(fixed_module_vars[0]), 0)

  def testDuplicateAssetCopy(self):
    export_path = os.path.join(self.get_temp_dir(), "assets-module")

    def module_with_duplicate_asset():
      vocabulary_file = self.create_vocab_file("tokens2.txt", ["1", "2", "3"])
      indices1 = tf.placeholder(dtype=tf.int64, name="indices1")
      indices2 = tf.placeholder(dtype=tf.int64, name="indices2")
      hub.add_signature(
          inputs={
              "indices_1": indices1,
              "indices_2": indices2,
          },
          outputs={
              "x": do_table_lookup(indices1, vocabulary_file),
              "y": do_table_lookup(indices2, vocabulary_file),
          })

    with tf.Graph().as_default():
      spec = hub.create_module_spec(module_with_duplicate_asset)
      module_a = hub.Module(spec)
      module_a({"indices_1": tf.constant([1, 2], dtype=tf.int64),
                "indices_2": tf.constant([1, 2], dtype=tf.int64)}, as_dict=True)
      with tf.Session() as sess:
        sess.run(tf.tables_initializer())
        module_a.export(export_path, sess)

 def testLoadModuleFromFuncDef(self):
   with tf.Session() as sess:
     v = tf.placeholder(tf.int64)
     spec = hub.create_module_spec(stateless_module_fn)
     m = hub.Module(spec)
     y = m(v)
     self.assertEqual(sess.run(y, feed_dict={v: 10}), 100)

  def test_module_export_vocab_on_custom_fs(self):
    root_dir = "file://%s" % self.get_temp_dir()
    export_dir = "%s_%s" % (root_dir, "export")
    tf.gfile.MakeDirs(export_dir)
    # Create a module with a vocab file located on a custom filesystem.
    vocab_dir = os.path.join(root_dir, "vocab_location")
    tf.gfile.MakeDirs(vocab_dir)
    vocab_filename = os.path.join(vocab_dir, "tokens.txt")
    tf_utils.atomic_write_string_to_file(vocab_filename, "one", False)

    def create_assets_module_fn():

      def assets_module_fn():
        indices = tf.placeholder(dtype=tf.int64, name="indices")
        table = tf.contrib.lookup.index_to_string_table_from_file(
            vocabulary_file=vocab_filename, default_value="UNKNOWN")
        outputs = table.lookup(indices)
        hub.add_signature(inputs=indices, outputs=outputs)

      return assets_module_fn

    with tf.Graph().as_default():
      assets_module_fn = create_assets_module_fn()
      spec = hub.create_module_spec(assets_module_fn)
      embedding_module = hub.Module(spec)
      with tf.Session() as sess:
        sess.run(tf.tables_initializer())
        embedding_module.export(export_dir, sess)

    module_files = tf.gfile.ListDirectory(export_dir)
    self.assertListEqual(
        ["assets", "saved_model.pb", "tfhub_module.pb", "variables"],
        sorted(module_files))
    module_files = tf.gfile.ListDirectory(os.path.join(export_dir, "assets"))
    self.assertListEqual(["tokens.txt"], module_files)

 def testUnusedInputModule(self):
   with tf.Session() as sess:
     v1 = tf.placeholder(tf.int64)
     v2 = tf.placeholder(tf.int64)
     spec = hub.create_module_spec(unused_input_module_fn)
     m = hub.Module(spec)
     out = m({"x": v1, "unused": v2})
     self.assertEqual(sess.run(out, feed_dict={v1: 10, v2: 4}), 100)

 def testVariables(self):
   spec = hub.create_module_spec(stateful_module_fn)
   m = hub.Module(spec, name="test")
   out = m()
   self.assertEqual(list(m.variable_map.keys()), ["var123"])
   self.assertEqual(m.variable_map["var123"].name, "test/var123:0")
   with tf.Session() as sess:
     sess.run(tf.global_variables_initializer())
     self.assertAllClose(sess.run(out), [1.0, 2.0, 3.0])

 def testMultipleOutputs(self):
   with tf.Session() as sess:
     spec = hub.create_module_spec(multiple_outputs_module_fn)
     m = hub.Module(spec)
     output = m(tf.constant([2.0]), as_dict=True)
     output1 = output["y"]
     output2 = output["z"]
     sess.run(tf.global_variables_initializer())
     self.assertAllClose(sess.run(output1), [6.0])
     self.assertAllClose(sess.run(output2), [18.0])

 def testConvertToTensor(self):
   spec = hub.create_module_spec(stateless_module_fn)
   with tf.Session() as sess:
     m = hub.Module(spec)
     y = m([10, 2])
     self.assertAllEqual(sess.run(y), [100, 4])
   with tf.Session() as sess:
     m = hub.Module(spec)
     with self.assertRaises(TypeError):
       m("hello")

 def testLargePartitionedVariables(self):
   spec = hub.create_module_spec(
       create_partitioned_variable_module_fn(partitions=25, shape=[600, 3]))
   m = hub.Module(spec, name="test")
   out = m()
   self.assertEqual(len(m.variable_map), 2)
   self.assertEqual(len(m.variable_map["partitioned_variable"]), 25)
   with tf.Session() as sess:
     sess.run(tf.global_variables_initializer())
     self.assertAllClose(sess.run(out), 2 * np.ones([600, 3]))

 def _testReluModule(self, module_fn):
   spec = hub.create_module_spec(module_fn)
   with tf.Graph().as_default():
     with tf.Session() as sess:
       x = tf.placeholder(dtype=tf.float32, name="x")
       relu_module = hub.Module(spec)
       y = relu_module(x)
       self.assertAllClose(sess.run(y, {x: 9.1}), 9.1)
       self.assertAllClose(sess.run(y, {x: -2.4}), 0.0)
       grad = tf.gradients([y], [x])
       self.assertAllClose(sess.run(grad, {x: 2}), [1.0])
       self.assertAllClose(sess.run(grad, {x: -2}), [0.0])

 def testWhileModule(self):
   spec = hub.create_module_spec(while_module_fn)
   with tf.Graph().as_default():
     with tf.Session() as sess:
       x = tf.placeholder(tf.float32)
       n = tf.placeholder(tf.int32)
       pow_module = hub.Module(spec)
       y = pow_module({"x": x, "n": n})
       self.assertAllClose(sess.run(y, {x: 9.1, n: 1}), 9.1)
       self.assertAllClose(sess.run(y, {x: 2.4, n: 2}), 5.76)
       grad = tf.gradients([y], [x])
       self.assertAllClose(sess.run(grad, {x: 2, n: 3}), [12.0])

  def testModuleWithMultipleSignatures(self):
    spec = hub.create_module_spec(multiple_signature_module_fn)
    module_a = hub.Module(spec, name="moduleA")
    in_tensor = tf.placeholder(dtype=tf.float32)
    out_tensor_a = module_a(in_tensor, signature="mul")
    out_tensor_b = module_a(out_tensor_a, signature="div")

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      in_values = [6, 3, 1]
      self.assertAllClose(
          sess.run(out_tensor_b, feed_dict={in_tensor: in_values}), in_values)

def export_module(path):
  spec = hub.create_module_spec(half_plus_two)

  with tf.Graph().as_default():
    module = hub.Module(spec)

    init_a = tf.assign(module.variable_map["a"], 0.5)
    init_b = tf.assign(module.variable_map["b"], 2.0)
    init_vars = tf.group([init_a, init_b])

    with tf.Session() as session:
      session.run(init_vars)
      module.export(path, session)

def export_as_tfhub_module(model_name,
                           hparams,
                           decode_hparams,
                           problem,
                           checkpoint_path,
                           export_dir):
  """Exports the last checkpoint from the directory as tfhub module.

  It creates the Module spec and signature (based on T2T problem information),
  which is later used to create and export the hub module.
  Module will be saved inside the ckpt_dir.

  Args:
    model_name: name of the model to be exported.
    hparams: T2T parameters, model graph will be based on them.
    decode_hparams: T2T parameters for decoding.
    problem: the name of the problem
    checkpoint_path: path to the checkpoint to be exported.
    export_dir: Directory to write the exported model to.
  """

  def hub_module_fn():
    """Creates the TF graph for the hub module."""
    model_fn = t2t_model.T2TModel.make_estimator_model_fn(
        model_name,
        hparams,
        decode_hparams=decode_hparams,
        use_tpu=FLAGS.use_tpu)
    features = problem.serving_input_fn(hparams).features

    # we must do a copy of the features, as the model_fn can add additional
    # entries there (like hyperparameter settings etc).
    original_features = features.copy()
    spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)

    hub.add_signature(
        inputs=original_features,
        outputs=spec.export_outputs["serving_default"].outputs)

  # TFHub doesn't support the following collections.
  drop_collections = [tf.GraphKeys.LOSSES,
                      tf.GraphKeys.SUMMARIES, tf.GraphKeys.LOCAL_VARIABLES]
  module_spec = hub.create_module_spec(
      hub_module_fn, drop_collections=drop_collections)
  # Loads the weights from the checkpoint using the model above
  # and saves it in the export_path.
  export_module_spec_with_checkpoint(
      module_spec,
      checkpoint_path=checkpoint_path,
      export_path=export_dir,
      scope_prefix="")

 def testExportedConsumerModelWorksIfItUsesHubModuleWithAssets(self):
   # 1. Create and export a module with assets.
   module_export_path = os.path.join(self.get_temp_dir(), "small-module")
   vocabulary_file = self.create_vocab_file("tokens.txt",
                                            ["emerson", "lake", "palmer"])
   assets_module_fn = create_assets_module_fn(vocabulary_file)
   spec = hub.create_module_spec(assets_module_fn)
   with tf.Graph().as_default():
     small_module = hub.Module(spec)
     with tf.Session() as sess:
       small_module.export(module_export_path, sess)
   # 2. Remove the original vocab file and move the module to another location.
   tf.gfile.Remove(vocabulary_file)
   inner_module_path = os.path.join(self.get_temp_dir(), "inner-module")
   tf.gfile.Rename(module_export_path, inner_module_path)
   del module_export_path
   # 3. Use the module in a consumer model (which is another module here).
   module_export_path = os.path.join(self.get_temp_dir(), "consumer-module")
   consumer_module_fn = create_consumer_module_fn(inner_module_path)
   spec = hub.create_module_spec(consumer_module_fn)
   with tf.Graph().as_default():
     consumer_module = hub.Module(spec)
     with tf.Session() as sess:
       consumer_module.export(module_export_path, sess)
   # 4. Delete the inner module on disk and move the consumer model to a final
   # location for serving.
   tf.gfile.DeleteRecursively(inner_module_path)
   module_serving_path = os.path.join(self.get_temp_dir(), "serving-module")
   tf.gfile.Rename(module_export_path, module_serving_path)
   # 5. Make sure the model can be served successfully.
   with tf.Graph().as_default():
     serving_module = hub.Module(module_serving_path)
     output = serving_module(tf.constant([1, 2], dtype=tf.int64))
     with tf.Session() as sess:
       sess.run(tf.tables_initializer())
       self.assertAllEqual(list(sess.run(output)), [b"lake", b"palmer"])

 def testVariableColocationPropagation(self):
   spec = hub.create_module_spec(stateful_module_fn_with_colocation)
   m = hub.Module(spec)
   u1 = tf.constant(1, name="u1")
   u2 = tf.constant(2, name="u2")
   with tf.colocate_with(u1), tf.colocate_with(u2):
     x = tf.constant(100.0, name="x")
   y = m(x)
   self.assertItemsEqual(y.op.colocation_groups(),
                         [tf.compat.as_bytes("loc:@module/var123"),
                          tf.compat.as_bytes("loc:@u1"),
                          tf.compat.as_bytes("loc:@u2")])
   with tf.Session() as sess:
     sess.run(tf.global_variables_initializer())
     self.assertEqual(sess.run(y), 101.0)