| 432 | |
| 433 | |
| 434 | def dot_generalized(a, b): |
| 435 | a = asarray(a) |
| 436 | if a.ndim >= 3: |
| 437 | if a.ndim == b.ndim: |
| 438 | # matrix x matrix |
| 439 | new_shape = a.shape[:-1] + b.shape[-1:] |
| 440 | elif a.ndim == b.ndim + 1: |
| 441 | # matrix x vector |
| 442 | new_shape = a.shape[:-1] |
| 443 | else: |
| 444 | raise ValueError("Not implemented...") |
| 445 | r = np.empty(new_shape, dtype=np.common_type(a, b)) |
| 446 | for c in itertools.product(*map(range, a.shape[:-2])): |
| 447 | r[c] = dot(a[c], b[c]) |
| 448 | return r |
| 449 | else: |
| 450 | return dot(a, b) |
| 451 | |
| 452 | |
| 453 | def identity_like_generalized(a): |