| 436 | } |
| 437 | |
| 438 | TensorLayout TensorLayout::broadcast(const TensorShape& tshape) const { |
| 439 | megdnn_throw_if( |
| 440 | !ndim || !tshape.ndim, tensor_reshape_error, |
| 441 | "broadcast involves empty tensor"); |
| 442 | |
| 443 | if (is_scalar()) { |
| 444 | TensorLayout result{dtype, format}; |
| 445 | result.ndim = tshape.ndim; |
| 446 | for (size_t i = 0; i < tshape.ndim; i++) { |
| 447 | result.shape[i] = tshape.shape[i]; |
| 448 | result.stride[i] = (tshape.shape[i] == 1); |
| 449 | } |
| 450 | return result; |
| 451 | } |
| 452 | |
| 453 | megdnn_throw_if( |
| 454 | tshape.ndim < ndim, tensor_reshape_error, |
| 455 | ssprintf( |
| 456 | "dimension for broadcast less than " |
| 457 | "dst_shape: src_shape=%s dst_shape=%s", |
| 458 | to_string().c_str(), tshape.to_string().c_str())); |
| 459 | TensorLayout result{dtype, format}; |
| 460 | for (size_t i = 0; i < tshape.ndim; ++i) { |
| 461 | int target_idx = tshape.ndim - i - 1; |
| 462 | int cur_idx = ndim - i - 1; |
| 463 | size_t cur_shape = (cur_idx >= 0 ? shape[cur_idx] : 1), |
| 464 | cur_stride = (cur_idx >= 0 ? stride[cur_idx] : 0); |
| 465 | if (tshape.shape[target_idx] != cur_shape) { |
| 466 | megdnn_throw_if( |
| 467 | cur_shape != 1 && cur_stride != 0, tensor_reshape_error, |
| 468 | ssprintf( |
| 469 | "broadcast on dim with shape not equal to 1: " |
| 470 | "src_shape=%s dst_shape=%s", |
| 471 | to_string().c_str(), tshape.to_string().c_str())); |
| 472 | result.shape[target_idx] = tshape.shape[target_idx]; |
| 473 | result.stride[target_idx] = 0; |
| 474 | } else { |
| 475 | result.shape[target_idx] = cur_shape; |
| 476 | result.stride[target_idx] = cur_stride; |
| 477 | } |
| 478 | } |
| 479 | result.ndim = tshape.ndim; |
| 480 | return result; |
| 481 | } |
| 482 | |
| 483 | bool TensorLayout::try_reshape(TensorLayout& result, const TensorShape& tshp) const { |
| 484 | megdnn_assert(tshp.ndim); |