# Buggy code
def __init__(
    self,
    optimizer: Optional[BlackboxOptimiser] = None,
    domain: Optional[str] = None,
    space: Optional[List[Dict]] = None,
    metric: Optional[str] = None,
    mode: Optional[str] = None,
    points_to_evaluate: Optional[List[List]] = None,
    evaluated_rewards: Optional[List] = None,
    **kwargs,
):
    assert dragonfly is not None, """dragonfly must be installed!
            You can install Dragonfly with the command:
            `pip install dragonfly-opt`."""
    if mode:
        assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."

    super(DragonflySearch, self).__init__(metric=metric, mode=mode, **kwargs)

    self._opt_arg = optimizer
    self._domain = domain
    self._space = space
    self._points_to_evaluate = points_to_evaluate
    self._evaluated_rewards = evaluated_rewards
    self._initial_points = []
    self._live_trial_mapping = {}

    self._opt = None
    if isinstance(optimizer, BlackboxOptimiser):
        if domain or space:
            raise ValueError(
                "If you pass an optimizer instance to dragonfly, do not "
                "pass a `domain` or `space`."
            )
        self._opt = optimizer
        self.init_dragonfly()
    elif self._space:
        self.setup_dragonfly()


# Traceback (TypeError)
Traceback (most recent call last):
(pid=82997)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/threading.py", line 916, in _bootstrap_inner
(pid=82997)     self.run()
(pid=82997)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 246, in run
(pid=82997)     raise e
(pid=82997)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 227, in run
(pid=82997)     self._entrypoint()
(pid=82997)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 290, in entrypoint
(pid=82997)     self._status_reporter.get_checkpoint())
(pid=82997)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 497, in _trainable_func
(pid=82997)     output = train_func(config)
(pid=82997)   File "test_ray.py", line 49, in trainer
(pid=82997)     keras.layers.Dense(config["neurons1"], input_shape=(784,), activation='relu', name='dense_1'),
(pid=82997)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/layers/core.py", line 1081, in __init__
(pid=82997)     self.units = int(units) if not isinstance(units, int) else units
(pid=82997) TypeError: int() argument must be a string, a bytes-like object or a number, not 'Integer'
------------------------------------------------------------
(pid=84324) Traceback (most recent call last):
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/threading.py", line 916, in _bootstrap_inner
(pid=84324)     self.run()
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 246, in run
(pid=84324)     raise e
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 227, in run
(pid=84324)     self._entrypoint()
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 290, in entrypoint
(pid=84324)     self._status_reporter.get_checkpoint())
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/ray/tune/function_runner.py", line 497, in _trainable_func
(pid=84324)     output = train_func(config)
(pid=84324)   File "test_ray.py", line 50, in trainer
(pid=84324)     keras.layers.Dense(10, activation='softmax', name='predictions'),
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
(pid=84324)     result = method(self, *args, **kwargs)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 116, in __init__
(pid=84324)     self.add(layer)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
(pid=84324)     result = method(self, *args, **kwargs)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 203, in add
(pid=84324)     output_tensor = layer(self.outputs[0])
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 773, in __call__
(pid=84324)     outputs = call_fn(cast_inputs, *args, **kwargs)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/layers/core.py", line 183, in call
(pid=84324)     lambda: array_ops.identity(inputs))
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/utils/tf_utils.py", line 59, in smart_cond
(pid=84324)     pred, true_fn=true_fn, false_fn=false_fn, name=name)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/smart_cond.py", line 59, in smart_cond
(pid=84324)     name=name)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
(pid=84324)     return func(*args, **kwargs)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/ops/control_flow_ops.py", line 1174, in cond
(pid=84324)     return cond_v2.cond_v2(pred, true_fn, false_fn, name)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/ops/cond_v2.py", line 83, in cond_v2
(pid=84324)     op_return_value=pred)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py", line 978, in func_graph_from_py_func
(pid=84324)     func_outputs = python_func(*func_args, **func_kwargs)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/keras/layers/core.py", line 179, in dropped_inputs
(pid=84324)     rate=self.rate)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
(pid=84324)     return func(*args, **kwargs)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/ops/nn_ops.py", line 4289, in dropout
(pid=84324)     return dropout_v2(x, rate, noise_shape=noise_shape, seed=seed, name=name)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/ops/nn_ops.py", line 4383, in dropout_v2
(pid=84324)     rate, dtype=x.dtype, name="rate")
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 1314, in convert_to_tensor
(pid=84324)     ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/constant_op.py", line 317, in _constant_tensor_conversion_function
(pid=84324)     return constant(v, dtype=dtype, name=name)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/constant_op.py", line 258, in constant
(pid=84324)     allow_broadcast=True)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/constant_op.py", line 296, in _constant_impl
(pid=84324)     allow_broadcast=allow_broadcast))
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/tensor_util.py", line 451, in make_tensor_proto
(pid=84324)     _AssertCompatible(values, dtype)
(pid=84324)   File "/home/gsukhorukov/.conda/envs/tf2/lib/python3.6/site-packages/tensorflow_core/python/framework/tensor_util.py", line 331, in _AssertCompatible
(pid=84324)     (dtype.name, repr(mismatch), type(mismatch).__name__))
(pid=84324) TypeError: Expected float32, got <ray.tune.sample.Categorical object at 0x7fe7d1835a90> of type 'Categorical' instead.

# Fixed code
def __init__(
    self,
    optimizer: Optional[BlackboxOptimiser] = None,
    domain: Optional[str] = None,
    space: Optional[Union[Dict, List[Dict]]] = None,
    metric: Optional[str] = None,
    mode: Optional[str] = None,
    points_to_evaluate: Optional[List[List]] = None,
    evaluated_rewards: Optional[List] = None,
    **kwargs,
):
    assert dragonfly is not None, """dragonfly must be installed!
            You can install Dragonfly with the command:
            `pip install dragonfly-opt`."""
    if mode:
        assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."

    super(DragonflySearch, self).__init__(metric=metric, mode=mode, **kwargs)

    self._opt_arg = optimizer
    self._domain = domain

    if isinstance(space, dict) and space:
        resolved_vars, domain_vars, grid_vars = parse_spec_vars(space)
        if domain_vars or grid_vars:
            logger.warning(UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self)))
            space = self.convert_search_space(space)

    self._space = space
    self._points_to_evaluate = points_to_evaluate
    self._evaluated_rewards = evaluated_rewards
    self._initial_points = []
    self._live_trial_mapping = {}

    self._opt = None
    if isinstance(optimizer, BlackboxOptimiser):
        if domain or space:
            raise ValueError(
                "If you pass an optimizer instance to dragonfly, do not "
                "pass a `domain` or `space`."
            )
        self._opt = optimizer
        self.init_dragonfly()
    elif self._space:
        self.setup_dragonfly()