# Buggy code
def _register_compressor(self, params, optimizer_params, compression_params):
    """Register compressor for BytePS

    params : mx.gluon.ParameterDict
    optimizer_params : dict
    compression_params : dict
    """
    intra_compressor = Compression.none
    if not compression_params:
        return intra_compressor

    if "fp16" in compression_params:
        intra_compressor = Compression.fp16

    if "compressor" not in compression_params:
        warnings.warn("Compressor is not defined")
        return intra_compressor

    check_list = ["compressor", "ef", "momentum"]

    for _, param in params.items():
        # generic
        for item in check_list:
            if compression_params.get(item):
                if isinstance(compression_params[item], str):
                    setattr(param, "byteps_%s_type" % item, compression_params[item])
                else:
                    raise TypeError("%s should be str" % item)

        # need parameter
        compressor = compression_params["compressor"]
        if compressor == "onebit":
            setattr(
                param,
                "byteps_compressor_onebit_scaling",
                str(compression_params.get("scaling", False)),
            )
        elif (
            compressor == "topk" or compressor == "randomk" or compressor == "multibit"
        ):
            # raise KeyError if 'k' is not found
            setattr(param, "byteps_compressor_k", compression_params["k"])

        if compression_params.get("momentum"):
            setattr(param, "byteps_momentum_mu", optimizer_params["momentum"])

    # change
    if compression_params.get("momentum"):
        # 1bit compressor use an additional momentum for weight decay
        if compressor == "onebit" and "wd" in optimizer_params:
            intra_compressor = Compression.wdmom(
                intra_compressor, optimizer_params["momentum"], optimizer_params["wd"]
            )
            del optimizer_params["wd"]

        del optimizer_params["momentum"]

    return intra_compressor


# Traceback (AttributeError)
$ DMLC_ROLE=worker DMLC_PS_ROOT_URI=12.12.10.12 DMLC_PS_ROOT_PORT=9000 DMLC_WORKER_ID=0 DMLC_NUM_WORKER=1 DMLC_NUM_SERVER=1 python launcher/launch.py python example/tensorflow/tensorflow_mnist.py
BytePS launching worker
INFO:tensorflow:Create CheckpointSaverHook.
Traceback (most recent call last):
File "example/tensorflow/tensorflow_mnist.py", line 160, in <module>
tf.app.run()
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/tensorflow/python/platform/app.py", line 125, in run
_sys.exit(main(argv))
File "example/tensorflow/tensorflow_mnist.py", line 152, in main
config=config) as mon_sess:
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 421, in MonitoredTrainingSession
stop_grace_period_secs=stop_grace_period_secs)
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 832, in __init__
stop_grace_period_secs=stop_grace_period_secs)
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/tensorflow/python/training/monitored_session.py", line 548, in __init__
h.begin()
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/byteps-0.1.0-py2.7-linux-x86_64.egg/byteps/tensorflow/__init__.py", line 107, in begin
self.bcast_op = broadcast_global_variables(self.root_rank)
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/byteps-0.1.0-py2.7-linux-x86_64.egg/byteps/tensorflow/__init__.py", line 66, in broadcast_global_variables
return broadcast_variables(tf.global_variables(), root_rank, scope)
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/byteps-0.1.0-py2.7-linux-x86_64.egg/byteps/tensorflow/__init__.py", line 78, in broadcast_variables
for var in variables])
File "/home/shuai/.conda/envs/byteps/lib/python2.7/site-packages/byteps-0.1.0-py2.7-linux-x86_64.egg/byteps/tensorflow/ops.py", line 116, in broadcast
scope = tf.compat.v1.get_default_graph().get_name_scope()
AttributeError: 'module' object has no attribute 'v1'
^C^C^C^C^C^C^C^C^C^CTerminated

# Fixed code
def _register_compressor(self, params, optimizer_params, compression_params):
    """Register compressor for BytePS

    params : mx.gluon.ParameterDict
    optimizer_params : dict
    compression_params : dict
    """
    intra_compressor = Compression.none
    if not compression_params:
        return intra_compressor

    if "fp16" in compression_params:
        intra_compressor = Compression.fp16

    if "compressor" not in compression_params:
        warnings.warn("Compressor is not defined")
        return intra_compressor

    check_list = ["compressor", "ef", "momentum"]

    for _, param in params.items():
        # generic
        for item in check_list:
            if compression_params.get(item):
                if isinstance(compression_params[item], str):
                    setattr(param, "byteps_%s_type" % item, compression_params[item])
                else:
                    raise TypeError("%s should be str" % item)

        # need parameter
        compressor = compression_params["compressor"]
        if compressor == "onebit":
            setattr(
                param,
                "byteps_compressor_onebit_scaling",
                str(compression_params.get("scaling", False)),
            )
        elif (
            compressor == "topk" or compressor == "randomk" or compressor == "multibit"
        ):
            # raise KeyError if 'k' is not found
            setattr(param, "byteps_compressor_k", compression_params["k"])

        if compression_params.get("momentum"):
            setattr(param, "byteps_momentum_mu", optimizer_params["momentum"])

        if compression_params.get("seed"):
            setattr(param, "byteps_seed", compression_params["seed"])

    # change
    if compression_params.get("momentum"):
        # 1bit compressor use an additional momentum for weight decay
        if compressor == "onebit" and "wd" in optimizer_params:
            intra_compressor = Compression.wdmom(
                intra_compressor, optimizer_params["momentum"], optimizer_params["wd"]
            )
            del optimizer_params["wd"]

        del optimizer_params["momentum"]

    return intra_compressor