MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / infer_output_attrs_fallible

Function infer_output_attrs_fallible

imperative/src/impl/ops/rng.cpp:875–894  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

873
874template <typename Op>
875std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible(
876 const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) {
877 bool success = inputs[0].layout.ndim != 0;
878 LogicalTensorDesc dest;
879 auto&& xxx_rng_def = def.cast_final_safe<Op>();
880 size_t nr_inp = inputs.size();
881 constexpr bool rng_with_shape = OpMeth<Op>::DnnOp::NR_INPUTS == 0;
882 if (rng_with_shape) {
883 mgb_assert(
884 nr_inp == 1, "%s expects 1 inputs; got %lu actually",
885 xxx_rng_def.dyn_typeinfo()->name, nr_inp);
886 }
887 dest.comp_node = inputs[0].comp_node;
888 if (success) {
889 dest.layout = _InferLayout<rng_with_shape>::do_infer(inputs[0], xxx_rng_def);
890 } else {
891 dest.layout = TensorLayout(inputs[0].layout.dtype);
892 }
893 return {{dest}, inputs[0].layout.ndim != 0};
894}
895
896template <>
897std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible<

Callers 7

TESTFunction · 0.50
check_rng_basicFunction · 0.50
dispatch_default_cpuMethod · 0.50
dispatch_kernelMethod · 0.50

Calls 2

TensorLayoutClass · 0.85
sizeMethod · 0.45

Tested by 5

TESTFunction · 0.40
check_rng_basicFunction · 0.40