Method compute_shapes

src/module.cpp:751–809

}

std::vector<shape> module::compute_shapes(const std::vector<shape>& inputs,
                                          compute_shapes_options options) const
{
    auto params = this->get_parameter_names();
    std::sort(params.begin(), params.end());
    std::unordered_map<instruction_ref, shape> ins_shapes;
    std::unordered_map<std::string, shape> adjusted_param_shapes;
    std::transform(inputs.begin(),
                   inputs.end(),
                   params.begin(),
                   std::inserter(adjusted_param_shapes, adjusted_param_shapes.end()),
                   [](auto ps, const auto& name) { return std::make_pair(name, ps); });
    for(auto ins : iterator_for(*this))
    {
        if(ins->name() == "@param")
        {
            ins_shapes[ins] =
                adjusted_param_shapes[any_cast<builtin::param>(ins->get_operator()).parameter];
            if(options.strict_type and ins->get_shape().type() != ins_shapes[ins].type())
            {
                MIGRAPHX_THROW(options.name + ": Mismatched type: expected " +
                               ins->get_shape().type_string() + " but passed " +
                               ins_shapes[ins].type_string());
            }
            if(options.strict_lens and ins->get_shape().lens() != ins_shapes[ins].lens())
            {
                MIGRAPHX_THROW(options.name + ": Mismatched lens: expected {" +
                               to_string_range(ins->get_shape().lens()) + "} but passed {" +
                               to_string_range(ins_shapes[ins].lens()) + "}");
            }
        }
        else if(ins->name() == "@literal")
        {
            if(not options.scalar_const_out_lens.empty() and ins->get_shape().scalar())
            {
                std::vector<std::size_t> strides(options.scalar_const_out_lens.size());
                ins_shapes[ins] =
                    shape{ins->get_shape().type(), options.scalar_const_out_lens, strides};
            }
            else
            {
                ins_shapes[ins] = ins->get_shape();
            }
        }
        else
        {
            std::vector<shape> input_shapes;
            input_shapes.resize(ins->inputs().size());
            std::transform(ins->inputs().begin(),
                           ins->inputs().end(),
                           input_shapes.begin(),
                           [&](auto in) { return ins_shapes.at(in); });
            if(ins->name() == "@return")
                return input_shapes;
            ins_shapes[ins] = ins->get_operator().compute_shape(input_shapes, ins->module_inputs());
        }
    }
    MIGRAPHX_THROW("No return found in the submodule");