How to use the sockeye.utils function in sockeye

def _get_later_major_version():
    release, major, minor = utils.parse_version(__version__)
    return "%s.%d.%s" % (release, int(major) + 1, minor)

def test_expand_requested_device_ids_exception(requested_device_ids, num_gpus_available):
    with pytest.raises(ValueError):
        utils._expand_requested_device_ids(requested_device_ids, num_gpus_available)

use_source_factor, perplexity_thresh, bleu_thresh):
    """Task: sort short sequences of digits"""
    with tmp_digits_dataset("test_seq_sort.", _TRAIN_LINE_COUNT, _TRAIN_LINE_COUNT_EMPTY, _LINE_MAX_LENGTH, _DEV_LINE_COUNT, _LINE_MAX_LENGTH,
                            _TEST_LINE_COUNT, _TEST_LINE_COUNT_EMPTY, _TEST_MAX_LENGTH,
                            sort_target=True, seed_train=_SEED_TRAIN_DATA, seed_dev=_SEED_DEV_DATA,
                            with_source_factors=use_source_factor) as data:
        data = check_train_translate(train_params=train_params,
                                     translate_params=translate_params,
                                     data=data,
                                     use_prepared_data=use_prepared_data,
                                     max_seq_len=_LINE_MAX_LENGTH + C.SPACE_FOR_XOS,
                                     compare_output=True,
                                     seed=seed)

        # get best validation perplexity
        metrics = sockeye.utils.read_metrics_file(os.path.join(data['model'], C.METRICS_NAME))
        perplexity = min(m[C.PERPLEXITY + '-val'] for m in metrics)

        # compute metrics
        bleu = sockeye.evaluate.raw_corpus_bleu(hypotheses=data['test_outputs'], references=data['test_targets'])
        chrf = sockeye.evaluate.raw_corpus_chrf(hypotheses=data['test_outputs'], references=data['test_targets'])
        bleu_restrict = sockeye.evaluate.raw_corpus_bleu(hypotheses=data['test_outputs_restricted'],
                                                         references=data['test_targets'])

        logger.info("test: %s", name)
        logger.info("perplexity=%f, bleu=%f, bleu_restrict=%f chrf=%f", perplexity, bleu, bleu_restrict, chrf)
        assert perplexity <= perplexity_thresh
        assert bleu >= bleu_thresh
        assert bleu_restrict >= bleu_thresh

data = utils.cast_conditionally(data, self.dtype)
        if self.config.dropout_prepost > 0.0:
            data = mx.sym.Dropout(data=data, p=self.config.dropout_prepost)

        # (batch_size * heads, 1, max_length)
        bias = mx.sym.expand_dims(transformer.get_variable_length_bias(lengths=data_length,
                                                                       max_length=seq_len,
                                                                       num_heads=self.config.attention_heads,
                                                                       fold_heads=True,
                                                                       name="%sbias" % self.prefix), axis=1)
        bias = utils.cast_conditionally(bias, self.dtype)
        for i, layer in enumerate(self.layers):
            # (batch_size, seq_len, config.model_size)
            data = layer(data, bias)
        data = self.final_process(data=data, prev=None)
        data = utils.uncast_conditionally(data, self.dtype)
        return data, data_length, seq_len

def extract(param_path: str,
            param_names: List[str],
            list_all: bool) -> Dict[str, np.ndarray]:
    """
    Extract specific parameters given their names.

    :param param_path: Path to the parameter file.
    :param param_names: Names of parameters to be extracted.
    :param list_all: List names of all available parameters.
    :return: Extracted parameter dictionary.
    """
    logger.info("Loading parameters from '%s'", param_path)
    arg_params, aux_params = utils.load_params(param_path)

    ext_params = {}  # type: Dict[str, np.ndarray]
    param_names = _extract(param_names, arg_params, ext_params)
    param_names = _extract(param_names, aux_params, ext_params)

    if len(param_names) > 0:
        logger.info("The following parameters were not found:")
        for name in param_names:
            logger.info("\t%s", name)
        logger.info("Check the following availabilities")
        list_all = True

    if list_all:
        if arg_params:
            logger.info("Available arg parameters:")
            for name in arg_params:

args.output = temp_dir.name
        args.max_updates = 0

    utils.seed_rngs(args.seed)

    check_arg_compatibility(args)
    output_folder = os.path.abspath(args.output)
    resume_training = check_resume(args, output_folder)

    setup_main_logger(file_logging=not args.no_logfile,
                      console=not args.quiet,
                      path=os.path.join(output_folder, C.LOG_NAME),
                      level=args.loglevel)
    if hasattr(args, "checkpoint_frequency"):
        logger.warning("'--checkpoint-frequency' is deprecated, and will be removed in the future.  Please use '--checkpoint-interval'")
    utils.log_basic_info(args)
    arguments.save_args(args, os.path.join(output_folder, C.ARGS_STATE_NAME))

    max_seq_len_source, max_seq_len_target = args.max_seq_len
    # The maximum length is the length before we add the BOS/EOS symbols
    max_seq_len_source = max_seq_len_source + C.SPACE_FOR_XOS
    max_seq_len_target = max_seq_len_target + C.SPACE_FOR_XOS
    logger.info("Adjusting maximum length to reserve space for a BOS/EOS marker. New maximum length: (%d, %d)",
                max_seq_len_source, max_seq_len_target)

    check_condition(args.length_task is not None or C.LENRATIO_MSE not in args.metrics,
                    "%s metrics requires enabling length ratio prediction with --length-task." % C.LENRATIO_MSE)

    with ExitStack() as exit_stack:
        context = utils.determine_context(device_ids=args.device_ids,
                                          use_cpu=args.use_cpu,
                                          disable_device_locking=args.disable_device_locking,

def log_linear_interpolation(predictions):
        log_probs = utils.average_arrays([p.log() for p in predictions])
        return -log_probs.log_softmax()  # pylint: disable=invalid-unary-operand-type

def _encode(self, F, data: mx.sym.Symbol, data_length: mx.sym.Symbol) -> mx.sym.Symbol:
        data = utils.cast_conditionally(F, data, self.dtype)
        if self.config.dropout_prepost > 0.0:
            data = F.Dropout(data=data, p=self.config.dropout_prepost)

        # (batch_size * heads, 1, seq_len)
        bias = F.expand_dims(self.valid_length_mask(data, data_length), axis=1)
        bias = utils.cast_conditionally(F, bias, self.dtype)
        for layer in self.layers:
            # (batch_size, seq_len, config.model_size)
            data = layer(data, bias)
        data = self.final_process(data, None)
        data = utils.uncast_conditionally(F, data, self.dtype)
        return data

self._top = Top1()  # type: mx.gluon.HybridBlock
            elif self.sample is not None:
                self._top = SampleK(k=self.beam_size,
                                    n=self.sample,
                                    max_batch_size=self.max_batch_size)  # type: mx.gluon.HybridBlock
            else:
                self._top = TopK(k=self.beam_size,
                                 vocab_size=len(self.vocab_target))  # type: mx.gluon.HybridBlock

            self._top.initialize(ctx=self.context)
            self._top.hybridize(static_alloc=True, static_shape=True)
        else:
            if self.skip_topk:
                self._top = utils.top1  # type: Callable
            else:
                self._top = partial(utils.topk, k=self.beam_size)  # type: Callable

        self._sort_by_index = SortByIndex()
        self._sort_by_index.initialize(ctx=self.context)
        self._sort_by_index.hybridize(static_alloc=True, static_shape=True)

        self._update_finished = NormalizeAndUpdateFinished(pad_id=C.PAD_ID,
                                                           eos_id=self.vocab_target[C.EOS_SYMBOL],
                                                           length_penalty_alpha=self.length_penalty.alpha,
                                                           length_penalty_beta=self.length_penalty.beta)
        self._update_finished.initialize(ctx=self.context)
        self._update_finished.hybridize(static_alloc=True, static_shape=True)

        self._prune_hyps = PruneHypotheses(threshold=self.beam_prune, beam_size=self.beam_size)
        self._prune_hyps.initialize(ctx=self.context)
        self._prune_hyps.hybridize(static_alloc=True, static_shape=True)

:param data: Input data.
        :param data_length: Vector with sequence lengths.
        :param seq_len: Maximum sequence length.
        :return: Encoded versions of input data data, data_length, seq_len.
        """
        data = utils.cast_conditionally(data, self.dtype)
        if self.config.dropout_prepost > 0.0:
            data = mx.sym.Dropout(data=data, p=self.config.dropout_prepost)

        # (batch_size * heads, 1, max_length)
        bias = mx.sym.expand_dims(transformer.get_variable_length_bias(lengths=data_length,
                                                                       max_length=seq_len,
                                                                       num_heads=self.config.attention_heads,
                                                                       fold_heads=True,
                                                                       name="%sbias" % self.prefix), axis=1)
        bias = utils.cast_conditionally(bias, self.dtype)
        for i, layer in enumerate(self.layers):
            # (batch_size, seq_len, config.model_size)
            data = layer(data, bias)
        data = self.final_process(data=data, prev=None)
        data = utils.uncast_conditionally(data, self.dtype)
        return data, data_length, seq_len