# Buggy code
def __init__(
    self,
    space: Optional[List[Dict]] = None,
    metric: Optional[str] = None,
    mode: Optional[str] = None,
    parameter_constraints: Optional[List] = None,
    outcome_constraints: Optional[List] = None,
    ax_client: Optional[AxClient] = None,
    use_early_stopped_trials: Optional[bool] = None,
    max_concurrent: Optional[int] = None,
):
    assert ax is not None, "Ax must be installed!"
    if mode:
        assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."

    super(AxSearch, self).__init__(
        metric=metric,
        mode=mode,
        max_concurrent=max_concurrent,
        use_early_stopped_trials=use_early_stopped_trials,
    )

    self._ax = ax_client
    self._space = space
    self._parameter_constraints = parameter_constraints
    self._outcome_constraints = outcome_constraints

    self.max_concurrent = max_concurrent

    self._objective_name = metric
    self._parameters = []
    self._live_trial_mapping = {}

    if self._ax or self._space:
        self.setup_experiment()


# 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,
    space: Optional[Union[Dict, List[Dict]]] = None,
    metric: Optional[str] = None,
    mode: Optional[str] = None,
    parameter_constraints: Optional[List] = None,
    outcome_constraints: Optional[List] = None,
    ax_client: Optional[AxClient] = None,
    use_early_stopped_trials: Optional[bool] = None,
    max_concurrent: Optional[int] = None,
):
    assert ax is not None, "Ax must be installed!"
    if mode:
        assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."

    super(AxSearch, self).__init__(
        metric=metric,
        mode=mode,
        max_concurrent=max_concurrent,
        use_early_stopped_trials=use_early_stopped_trials,
    )

    self._ax = ax_client

    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._parameter_constraints = parameter_constraints
    self._outcome_constraints = outcome_constraints

    self.max_concurrent = max_concurrent

    self._objective_name = metric
    self._parameters = []
    self._live_trial_mapping = {}

    if self._ax or self._space:
        self.setup_experiment()