| 3355 | } |
| 3356 | |
| 3357 | array kron(const array& a, const array& b, StreamOrDevice s /* = {} */) { |
| 3358 | if (a.size() == 0 || b.size() == 0) { |
| 3359 | throw std::invalid_argument("[kron] Input arrays cannot be empty."); |
| 3360 | } |
| 3361 | |
| 3362 | int ndim = std::max(a.ndim(), b.ndim()); |
| 3363 | Shape a_shape(2 * ndim, 1); |
| 3364 | Shape b_shape(2 * ndim, 1); |
| 3365 | Shape out_shape(ndim, 1); |
| 3366 | |
| 3367 | for (int i = ndim - 1, j = a.ndim() - 1; j >= 0; j--, i--) { |
| 3368 | a_shape[2 * i] = a.shape(j); |
| 3369 | out_shape[i] *= a.shape(j); |
| 3370 | } |
| 3371 | for (int i = ndim - 1, j = b.ndim() - 1; j >= 0; j--, i--) { |
| 3372 | b_shape[2 * i + 1] = b.shape(j); |
| 3373 | out_shape[i] *= b.shape(j); |
| 3374 | } |
| 3375 | |
| 3376 | return reshape( |
| 3377 | multiply( |
| 3378 | reshape(a, std::move(a_shape), s), |
| 3379 | reshape(b, std::move(b_shape), s), |
| 3380 | s), |
| 3381 | std::move(out_shape), |
| 3382 | s); |
| 3383 | } |
| 3384 | |
| 3385 | array take( |
| 3386 | const array& a, |