Function infer_output_attrs_fallible

imperative/src/impl/ops/reduce.cpp:175–227

}

std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible(
        const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) {
    auto&& op_def = def.cast_final_safe<Reduce>();
    auto axis = op_def.param().axis;
    auto keepdim = op_def.keepdim;

    mgb_assert(inputs.size() > 0);
    auto&& comp_node = inputs[0].comp_node;
    auto&& input_layout = inputs[0].layout;

    if (inputs.size() == 2) {
        // fallback to proxy_graph, matters on backward
        auto [output_descs, validated] =
                proxy_graph_detail::infer_output_attrs_fallible(def, inputs);
        if (!inputs[1].value.empty()) {
            cg::copy_tensor_value_to_shape(output_descs[0].layout, inputs[1].value);
            output_descs[0].layout.init_contiguous_stride();
        }
        return {output_descs, validated};
    }

    mgb_assert(inputs.size() == 1);

    if (axis == INT_MAX) {
        // reduce to scalar
        // ignore keepdim because ndim is 1
        auto&& dtype = input_layout.dtype;
        auto&& format = input_layout.format;
        auto output_layout = TensorLayout{{1}, dtype, format};
        return {{{output_layout, comp_node}}, true};
    }

    if (input_layout.ndim == 0) {
        // shape incomplete
        return {{{TensorLayout(input_layout.dtype, input_layout.format), comp_node}},
                false};
    }

    if (axis < 0) {
        axis = input_layout.ndim + axis;
    }
    mgb_assert(axis >= 0 && axis < input_layout.ndim);

    TensorLayout output_layout = input_layout;
    bool remove_axis = (!keepdim) && input_layout.ndim > 1;
    if (remove_axis) {
        output_layout.remove_axis_inplace(axis);
    } else {
        output_layout.shape[axis] = 1;
    }
    output_layout.init_contiguous_stride();
    return {{{output_layout, comp_node}}, true};
}

SmallVector<VarNode::LayoutConstraintCallback> get_input_layout_constraint(
        const OpDef& def, const SmallVector<TensorPtr>& inputs) {