def compute_gather(params, indices, axis, batch_dims):
    """
    This utility function computes the gather operation with batch_dims supported.
    """
    def compute_gather_helper(params, indices, axis):
        scalar_indices = isinstance(indices, numbers.Integral)
        if scalar_indices:
            res = np.take(params, [indices], axis)
            res2 = np.squeeze(res, axis=axis)
            if isinstance(res2, np.ndarray) and len(res2.shape) == 0:
                # The `res2` is a numpy 0-d array (after doing np.squeeze on a 1-d array).
                # For 0-d array in numpy, we need to extract the scalar value by first converting
                # it back to 1-d array.
                # Notice that .item() doesn't work because it returns a built-in type instead of
                # np.generic type, which will fail the downstream var value setter.
                return np.atleast_1d(res2)[0]
            return res2
        return np.take(params, indices, axis)

    if batch_dims == 0:
        return compute_gather_helper(params, indices, axis)

    params_shape = params.shape
    indices_shape = indices.shape
    batch_shape = params_shape[:batch_dims]

    params_new_shape = [np.prod(batch_shape)] + list(params_shape[batch_dims:])
    indices_new_shape = [np.prod(batch_shape)] + list(indices_shape[batch_dims:])
    params_reshape = np.reshape(params, params_new_shape)
    indices_reshape = np.reshape(indices, indices_new_shape)

    res = []
    for p, i in zip(params_reshape, indices_reshape):
        res.append(compute_gather_helper(p, i, axis - batch_dims))
    res = np.stack(res)
    res_new_shape = tuple(batch_shape) + tuple(res.shape[1:])
    return np.reshape(res, res_new_shape)

    This utility function computes the gather operation with batch_dims supported.
    
coremltools/converters/mil/mil/ops/defs/_utils.py