Compute D = beta * C + alpha * (A @ B) */
| 5464 | |
| 5465 | /** Compute D = beta * C + alpha * (A @ B) */ |
| 5466 | array addmm( |
| 5467 | array c, |
| 5468 | array a, |
| 5469 | array b, |
| 5470 | const float& alpha /* = 1.f */, |
| 5471 | const float& beta /* = 1.f */, |
| 5472 | StreamOrDevice s /* = {} */) { |
| 5473 | int in_a_ndim = a.ndim(); |
| 5474 | int in_b_ndim = b.ndim(); |
| 5475 | |
| 5476 | if (a.ndim() == 0 || b.ndim() == 0) { |
| 5477 | throw std::invalid_argument( |
| 5478 | "[addmm] Got 0 dimension input. Inputs must " |
| 5479 | "have at least one dimension."); |
| 5480 | } |
| 5481 | |
| 5482 | // Type promotion |
| 5483 | auto out_type = result_type(a, b, c); |
| 5484 | |
| 5485 | if (out_type == complex64) { |
| 5486 | return add( |
| 5487 | multiply(matmul(a, b, s), array(alpha), s), |
| 5488 | multiply(array(beta), c, s), |
| 5489 | s); |
| 5490 | } |
| 5491 | |
| 5492 | if (a.ndim() == 1) { |
| 5493 | // Insert a singleton dim in the beginning |
| 5494 | a = expand_dims(a, 0, s); |
| 5495 | } |
| 5496 | if (b.ndim() == 1) { |
| 5497 | // Insert a singleton dim at the end |
| 5498 | b = expand_dims(b, 1, s); |
| 5499 | } |
| 5500 | |
| 5501 | if (a.shape(-1) != b.shape(-2)) { |
| 5502 | std::ostringstream msg; |
| 5503 | msg << "[addmm] Last dimension of first input with shape " << a.shape() |
| 5504 | << " must match second to last dimension of" |
| 5505 | << " second input with shape " << b.shape() << "."; |
| 5506 | throw std::invalid_argument(msg.str()); |
| 5507 | } |
| 5508 | |
| 5509 | if (!issubdtype(out_type, floating)) { |
| 5510 | std::ostringstream msg; |
| 5511 | msg << "[addmm] Only real floating point types are supported but " |
| 5512 | << c.dtype() << ", " << a.dtype() << " and " << b.dtype() |
| 5513 | << " were provided which results in " << out_type |
| 5514 | << ", which is not a real floating point type."; |
| 5515 | throw std::invalid_argument(msg.str()); |
| 5516 | } |
| 5517 | |
| 5518 | a = astype(a, out_type, s); |
| 5519 | b = astype(b, out_type, s); |
| 5520 | c = astype(c, out_type, s); |
| 5521 | |
| 5522 | // We can batch the multiplication by reshaping a |
| 5523 | if (a.ndim() > 2 && b.ndim() == 2 && c.ndim() <= 1) { |
no test coverage detected