| 527 | |
| 528 | template <typename Op> |
| 529 | SmallVector<LogicalTensorDesc> infer_output_attrs( |
| 530 | const OpDef& op, const SmallVector<TensorPtr>& inputs) { |
| 531 | LogicalTensorDesc dest; |
| 532 | auto&& rng = op.cast_final_safe<Op>(); |
| 533 | auto handle = rng.handle; |
| 534 | if (handle) { |
| 535 | dest.comp_node = RNGDnnOpManager::get_comp_node(handle); |
| 536 | } else { |
| 537 | dest.comp_node = inputs[0]->comp_node(); |
| 538 | } |
| 539 | constexpr bool rng_with_shape = OpMeth<Op>::DnnOp::NR_INPUTS == 0; |
| 540 | if (!rng_with_shape) { |
| 541 | for (int i = 0; i < inputs.size(); ++i) { |
| 542 | mgb_assert( |
| 543 | inputs[i]->comp_node() == dest.comp_node, |
| 544 | "%s expects the device of inputs[%d] to be same as the device of " |
| 545 | "handle; " |
| 546 | "got %s and %s actually", |
| 547 | rng.dyn_typeinfo()->name, i, |
| 548 | inputs[i]->comp_node().to_string().c_str(), |
| 549 | dest.comp_node.to_string().c_str()); |
| 550 | } |
| 551 | } |
| 552 | dest.layout = _InferLayout<rng_with_shape>::do_infer(inputs[0], rng); |
| 553 | return {dest}; |
| 554 | } |
| 555 | |
| 556 | template <> |
| 557 | SmallVector<LogicalTensorDesc> infer_output_attrs<ShuffleRNG>( |
nothing calls this directly
no test coverage detected