(i)
| 95 | The product ab |
| 96 | """ |
| 97 | def row_mul(i): |
| 98 | return torch.stack( |
| 99 | [ |
| 100 | a[..., i, 0] * b[..., 0, 0] |
| 101 | + a[..., i, 1] * b[..., 1, 0] |
| 102 | + a[..., i, 2] * b[..., 2, 0], |
| 103 | a[..., i, 0] * b[..., 0, 1] |
| 104 | + a[..., i, 1] * b[..., 1, 1] |
| 105 | + a[..., i, 2] * b[..., 2, 1], |
| 106 | a[..., i, 0] * b[..., 0, 2] |
| 107 | + a[..., i, 1] * b[..., 1, 2] |
| 108 | + a[..., i, 2] * b[..., 2, 2], |
| 109 | ], |
| 110 | dim=-1, |
| 111 | ) |
| 112 | |
| 113 | return torch.stack( |
| 114 | [ |