r"""Computes the linear optimal transport loss given a batch cost matrices and transport plans. .. math:: L(T, M)_b = \langle T_b, M_b \rangle_F Parameters ---------- M : array-like, shape (B, ns, nt) Cost matrix T : array-like, shape (B, ns, nt) Trans
(M, T, nx=None)
| 121 | |
| 122 | |
| 123 | def loss_linear_batch(M, T, nx=None): |
| 124 | r"""Computes the linear optimal transport loss given a batch cost matrices and transport plans. |
| 125 | |
| 126 | .. math:: |
| 127 | |
| 128 | L(T, M)_b = \langle T_b, M_b \rangle_F |
| 129 | |
| 130 | Parameters |
| 131 | ---------- |
| 132 | M : array-like, shape (B, ns, nt) |
| 133 | Cost matrix |
| 134 | T : array-like, shape (B, ns, nt) |
| 135 | Transport plan |
| 136 | Returns |
| 137 | ------- |
| 138 | loss : array-like, shape (B,) |
| 139 | Loss value for each batch element |
| 140 | See Also |
| 141 | -------- |
| 142 | ot.batch.dist_batch : batched cost matrix computation for computing M. |
| 143 | ot.batch.solve_batch : solver for computing the optimal T. |
| 144 | """ |
| 145 | if nx is None: |
| 146 | nx = get_backend(M, T) |
| 147 | return nx.sum(M * T, axis=(1, 2)) |
| 148 | |
| 149 | |
| 150 | def loss_linear_samples_batch(X, Y, T, metric="l2"): |