def METHOD_NAME(
    model, first_available_dim, temperature, rng_key, *args, **kwargs
):
    if temperature == 0:
        sum_op, prod_op = funsor.ops.max, funsor.ops.add
        approx = funsor.approximations.argmax_approximate
    elif temperature == 1:
        sum_op, prod_op = funsor.ops.logaddexp, funsor.ops.add
        rng_key, sub_key = random.split(rng_key)
        approx = funsor.montecarlo.MonteCarlo(rng_key=sub_key)
    else:
        raise ValueError("temperature must be 0 (map) or 1 (sample) for now")
    if first_available_dim is None:
        with block():
            model_trace = trace(seed(model, rng_key)).get_trace(*args, **kwargs)
        first_available_dim = -_guess_max_plate_nesting(model_trace) - 1
    with funsor.adjoint.AdjointTape() as tape:
        with block(), enum(first_available_dim=first_available_dim):
            log_prob, model_tr, log_measures = _enum_log_density(
                model, args, kwargs, {}, sum_op, prod_op
            )
    with approx:
        approx_factors = tape.adjoint(sum_op, prod_op, log_prob)
    # construct a result trace to replay against the model
    sample_tr = model_tr.copy()
    for name, node in sample_tr.items():
        if node["type"] != "sample":
            continue
        if node["infer"].get("enumerate") == "parallel":
            log_measure = approx_factors[log_measures[name]]
            value = _get_support_value(log_measure, name)
            node["value"] = funsor.to_data(
                value, name_to_dim=node["infer"]["name_to_dim"]
            )
    data = {
        name: site["value"]
        for name, site in sample_tr.items()
        if site["type"] == "sample"
    }
    # concatenate _PREV_foo to foo
    time_vars = defaultdict(list)
    for name in data:
        if name.startswith("_PREV_"):
            root_name = _shift_name(name, -_get_shift(name))
            time_vars[root_name].append(name)
    for name in time_vars:
        if name in data:
            time_vars[name].append(name)
        time_vars[name] = sorted(time_vars[name], key=len, reverse=True)
    for root_name, vars in time_vars.items():
        prototype_shape = model_trace[root_name]["value"].shape
        values = [data.pop(name) for name in vars]
        if len(values) == 1:
            data[root_name] = values[0].reshape(prototype_shape)
        else:
            assert len(prototype_shape) >= 1
            values = [v.reshape((-1,) + prototype_shape[1:]) for v in values]
            data[root_name] = jnp.concatenate(values)
    return data