# Path: disentanglement_lib/utils/convolute_hub.py
# def convolute_and_save(module_path, signature, export_path, transform_fn,
#                        transform_checkpoint_path, new_signature=None):
#   def module_fn():
# def save_numpy_arrays_to_checkpoint(checkpoint_path, **dict_with_arrays):
# def _placeholders_from_module(tfhub_module, signature):


import os
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from disentanglement_lib.utils import convolute_hub
    random_state = np.random.RandomState(0)
    data = random_state.normal(size=(5, 10))
    variable1 = random_state.normal(size=(10, 6))
    variable2 = random_state.normal(size=(6, 2))

    # Save variables to checkpoint.
    checkpoint_path = os.path.join(self.get_temp_dir(), "checkpoint.ckpt")
    convolute_hub.save_numpy_arrays_to_checkpoint(
        checkpoint_path, variable1=variable1, variable2=variable2)

    # Save a TFHub module that we will convolute.
    module_path = os.path.join(self.get_temp_dir(), "module_path")
    def module_fn():
      tensor = tf.placeholder(tf.float64, shape=(None, 10))
      variable1 = tf.get_variable("variable1", shape=(10, 6), dtype=tf.float64)
      output = tf.matmul(tensor, variable1)
      hub.add_signature(
          name="multiplication1",
          inputs={"tensor": tensor},
          outputs={"tensor": output})
    spec = hub.create_module_spec(module_fn)
    spec.export(module_path, checkpoint_path=checkpoint_path)

    # Function used for the convolution.
    def _operation2(tensor):
      variable2 = tf.get_variable("variable2", shape=(6, 2), dtype=tf.float64)
      return dict(tensor=tf.matmul(tensor, variable2))

    # Save the convolution as a new TFHub module
    module_path_new = os.path.join(self.get_temp_dir(), "module_path_new")

# Path: disentanglement_lib/utils/convolute_hub.py
# def convolute_and_save(module_path, signature, export_path, transform_fn,
#                        transform_checkpoint_path, new_signature=None):
#   def module_fn():
# def save_numpy_arrays_to_checkpoint(checkpoint_path, **dict_with_arrays):
# def _placeholders_from_module(tfhub_module, signature):


# Path: disentanglement_lib/utils/convolute_hub_test.py
import os
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
from disentanglement_lib.utils import convolute_hub

    random_state = np.random.RandomState(0)
    data = random_state.normal(size=(5, 10))
    variable1 = random_state.normal(size=(10, 6))
    variable2 = random_state.normal(size=(6, 2))

    # Save variables to checkpoint.
    checkpoint_path = os.path.join(self.get_temp_dir(), "checkpoint.ckpt")
    convolute_hub.save_numpy_arrays_to_checkpoint(
        checkpoint_path, variable1=variable1, variable2=variable2)

    # Save a TFHub module that we will convolute.
    module_path = os.path.join(self.get_temp_dir(), "module_path")
    def module_fn():
      tensor = tf.placeholder(tf.float64, shape=(None, 10))
      variable1 = tf.get_variable("variable1", shape=(10, 6), dtype=tf.float64)
      output = tf.matmul(tensor, variable1)
      hub.add_signature(
          name="multiplication1",
          inputs={"tensor": tensor},
          outputs={"tensor": output})
    spec = hub.create_module_spec(module_fn)
    spec.export(module_path, checkpoint_path=checkpoint_path)

    # Function used for the convolution.
    def _operation2(tensor):
      variable2 = tf.get_variable("variable2", shape=(6, 2), dtype=tf.float64)
      return dict(tensor=tf.matmul(tensor, variable2))

    # Save the convolution as a new TFHub module
    module_path_new = os.path.join(self.get_temp_dir(), "module_path_new")

    convolute_hub.convolute_and_save(module_path, "multiplication1",