r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. Parameters ---------- M : array-like, shape (B, ns, nt) Cost matrix reg : float Regularization parameter for entropic regularization a : array-like, shape (B, ns) Sour
(
M,
reg,
a=None,
b=None,
max_iter=1000,
tol=1e-5,
solver="log_sinkhorn",
reg_type="entropy",
grad="envelope",
)
| 234 | |
| 235 | |
| 236 | def solve_batch( |
| 237 | M, |
| 238 | reg, |
| 239 | a=None, |
| 240 | b=None, |
| 241 | max_iter=1000, |
| 242 | tol=1e-5, |
| 243 | solver="log_sinkhorn", |
| 244 | reg_type="entropy", |
| 245 | grad="envelope", |
| 246 | ): |
| 247 | r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. |
| 248 | |
| 249 | Parameters |
| 250 | ---------- |
| 251 | M : array-like, shape (B, ns, nt) |
| 252 | Cost matrix |
| 253 | reg : float |
| 254 | Regularization parameter for entropic regularization |
| 255 | a : array-like, shape (B, ns) |
| 256 | Source distribution (optional). If None, uniform distribution is used. |
| 257 | b : array-like, shape (B, nt) |
| 258 | Target distribution (optional). If None, uniform distribution is used. |
| 259 | max_iter : int |
| 260 | Maximum number of iterations |
| 261 | tol : float |
| 262 | Tolerance for convergence |
| 263 | solver: str |
| 264 | Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. |
| 265 | reg_type : str, optional |
| 266 | Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". |
| 267 | grad : str, optional |
| 268 | Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for |
| 269 | Sinkhorn solver. By default 'autodiff' provides gradients wrt all |
| 270 | outputs (`plan, value, value_linear`) but with important memory cost. |
| 271 | 'envelope' provides gradients only for `value` and and other outputs are |
| 272 | detached. This is useful for memory saving when only the value is needed. 'last_step' provides |
| 273 | gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. |
| 274 | 'detach' does not compute the gradients for the Sinkhorn solver. |
| 275 | |
| 276 | Returns |
| 277 | ------- |
| 278 | res : OTResult() |
| 279 | Result of the optimization problem. The information can be obtained as follows: |
| 280 | |
| 281 | - res.plan : OT plan :math:`\mathbf{T}` |
| 282 | - res.potentials : OT dual potentials |
| 283 | - res.value : Optimal value of the optimization problem |
| 284 | - res.value_linear : Linear OT loss with the optimal OT plan |
| 285 | |
| 286 | See :any:`OTResult` for more information. |
| 287 | |
| 288 | Examples |
| 289 | -------- |
| 290 | >>> import numpy as np |
| 291 | >>> from ot.batch import solve_batch, dist_batch |
| 292 | >>> X = np.random.randn(5, 10, 3) # 5 batches of 10 samples in 3D |
| 293 | >>> Y = np.random.randn(5, 15, 3) # 5 batches of 15 samples in 3D |