# Buggy code
def balanced_accuracy(solution, prediction):
    y_type, solution, prediction = _check_targets(solution, prediction)

    if y_type not in ["binary", "multiclass", "multilabel-indicator"]:
        raise ValueError(f"{y_type} is not supported")

    if y_type == "binary":
        # Do not transform into any multiclass representation
        pass

    elif y_type == "multiclass":
        # Need to create a multiclass solution and a multiclass predictions
        max_class = int(np.max((np.max(solution), np.max(prediction))))
        solution_binary = np.zeros((len(solution), max_class + 1))
        prediction_binary = np.zeros((len(prediction), max_class + 1))
        for i in range(len(solution)):
            solution_binary[i, int(solution[i])] = 1
            prediction_binary[i, int(prediction[i])] = 1
        solution = solution_binary
        prediction = prediction_binary

    elif y_type == "multilabel-indicator":
        solution = solution.toarray()
        prediction = prediction.toarray()
    else:
        raise NotImplementedError(f"bac_metric does not support task type {y_type}")

    fn = np.sum(np.multiply(solution, (1 - prediction)), axis=0, dtype=float)
    tp = np.sum(np.multiply(solution, prediction), axis=0, dtype=float)
    # Bounding to avoid division by 0
    eps = 1e-15
    tp = np.maximum(eps, tp)
    pos_num = np.maximum(eps, tp + fn)
    tpr = tp / pos_num  # true positive rate (sensitivity)

    if y_type in ("binary", "multilabel-indicator"):
        tn = np.sum(np.multiply((1 - solution), (1 - prediction)), axis=0, dtype=float)
        fp = np.sum(np.multiply((1 - solution), prediction), axis=0, dtype=float)
        tn = np.maximum(eps, tn)
        neg_num = np.maximum(eps, tn + fp)
        tnr = tn / neg_num  # true negative rate (specificity)
        bac = 0.5 * (tpr + tnr)
    elif y_type == "multiclass":
        bac = tpr
    else:
        raise ValueError(y_type)

    return np.mean(bac)  # average over all classes


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-28-ddcca1d23e69> in <module>
----> 1 ag.utils.tabular.metrics.balanced_accuracy(y_test, y_pred)

~/Desktop/github-aws/autogluon-public/autogluon/utils/tabular/metrics/__init__.py in __call__(self, y_true, y_pred, sample_weight)
87         else:
88             return self._sign * self._score_func(y_true, y_pred,
---> 89                                                  **self._kwargs)
90
91

~/Desktop/github-aws/autogluon-public/autogluon/utils/tabular/metrics/classification_metrics.py in balanced_accuracy(solution, prediction)
19     elif y_type == 'multiclass':
20         # Need to create a multiclass solution and a multiclass predictions
---> 21         max_class = int(np.max((np.max(solution), np.max(prediction))))
22         solution_binary = np.zeros((len(solution), max_class + 1))
23         prediction_binary = np.zeros((len(prediction), max_class + 1))

<__array_function__ internals> in amax(*args, **kwargs)

~/Desktop/github-aws/ghaws/lib/python3.7/site-packages/numpy/core/fromnumeric.py in amax(a, axis, out, keepdims, initial, where)
2666     """
2667     return _wrapreduction(a, np.maximum, 'max', axis, None, out,
-> 2668                           keepdims=keepdims, initial=initial, where=where)
2669
2670

~/Desktop/github-aws/ghaws/lib/python3.7/site-packages/numpy/core/fromnumeric.py in _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs)
88                 return reduction(axis=axis, out=out, **passkwargs)
89
---> 90     return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
91
92

TypeError: cannot perform reduce with flexible type

# Fixed code
def balanced_accuracy(solution, prediction):
    y_type, solution, prediction = _check_targets(solution, prediction)

    if y_type not in ["binary", "multiclass", "multilabel-indicator"]:
        raise ValueError(f"{y_type} is not supported")

    if y_type == "binary":
        # Do not transform into any multiclass representation
        pass

    elif y_type == "multiclass":
        n = len(solution)
        unique_sol, encoded_sol = np.unique(solution, return_inverse=True)
        unique_pred, encoded_pred = np.unique(prediction, return_inverse=True)
        classes = np.unique(np.concatenate((unique_sol, unique_pred)))
        map_sol = np.array([np.where(classes == c)[0][0] for c in unique_sol])
        map_pred = np.array([np.where(classes == c)[0][0] for c in unique_pred])
        # one hot encoding
        sol_ohe = np.zeros((n, len(classes)))
        pred_ohe = np.zeros((n, len(classes)))
        sol_ohe[np.arange(n), map_sol[encoded_sol]] = 1
        pred_ohe[np.arange(n), map_pred[encoded_pred]] = 1
        solution = sol_ohe
        prediction = pred_ohe

    elif y_type == "multilabel-indicator":
        solution = solution.toarray()
        prediction = prediction.toarray()
    else:
        raise NotImplementedError(f"bac_metric does not support task type {y_type}")

    fn = np.sum(np.multiply(solution, (1 - prediction)), axis=0, dtype=float)
    tp = np.sum(np.multiply(solution, prediction), axis=0, dtype=float)
    # Bounding to avoid division by 0
    eps = 1e-15
    tp = np.maximum(eps, tp)
    pos_num = np.maximum(eps, tp + fn)
    tpr = tp / pos_num  # true positive rate (sensitivity)

    if y_type in ("binary", "multilabel-indicator"):
        tn = np.sum(np.multiply((1 - solution), (1 - prediction)), axis=0, dtype=float)
        fp = np.sum(np.multiply((1 - solution), prediction), axis=0, dtype=float)
        tn = np.maximum(eps, tn)
        neg_num = np.maximum(eps, tn + fp)
        tnr = tn / neg_num  # true negative rate (specificity)
        bac = 0.5 * (tpr + tnr)
    elif y_type == "multiclass":
        bac = tpr
    else:
        raise ValueError(y_type)

    return np.mean(bac)  # average over all classes