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 metric : str, optional 'sqeuclidea
(
X_a,
X_b,
reg,
a=None,
b=None,
metric="sqeuclidean",
p=2,
max_iter=1000,
tol=1e-5,
solver="log_sinkhorn",
reg_type="entropy",
grad="envelope",
)
| 358 | |
| 359 | |
| 360 | def solve_sample_batch( |
| 361 | X_a, |
| 362 | X_b, |
| 363 | reg, |
| 364 | a=None, |
| 365 | b=None, |
| 366 | metric="sqeuclidean", |
| 367 | p=2, |
| 368 | max_iter=1000, |
| 369 | tol=1e-5, |
| 370 | solver="log_sinkhorn", |
| 371 | reg_type="entropy", |
| 372 | grad="envelope", |
| 373 | ): |
| 374 | r"""Batched version of ot.solve, use it to solve many entropic OT problems in parallel. |
| 375 | |
| 376 | Parameters |
| 377 | ---------- |
| 378 | M : array-like, shape (B, ns, nt) |
| 379 | Cost matrix |
| 380 | reg : float |
| 381 | Regularization parameter for entropic regularization |
| 382 | metric : str, optional |
| 383 | 'sqeuclidean', 'euclidean', 'minkowski' or 'kl' |
| 384 | p : float, optional |
| 385 | p-norm for the Minkowski metrics. Default value is 2. |
| 386 | a : array-like, shape (B, ns) |
| 387 | Source distribution (optional). If None, uniform distribution is used. |
| 388 | b : array-like, shape (B, nt) |
| 389 | Target distribution (optional). If None, uniform distribution is used. |
| 390 | max_iter : int |
| 391 | Maximum number of iterations |
| 392 | tol : float |
| 393 | Tolerance for convergence |
| 394 | solver: str |
| 395 | Solver to use, either 'log_sinkhorn' or 'sinkhorn'. Default is "log_sinkhorn" which is more stable. |
| 396 | reg_type : str, optional |
| 397 | Type of regularization :math:`R` either "KL", or "entropy". Default is "entropy". |
| 398 | grad : str, optional |
| 399 | Type of gradient computation, either or 'autodiff', 'envelope' or 'last_step' used only for |
| 400 | Sinkhorn solver. By default 'autodiff' provides gradients wrt all |
| 401 | outputs (`plan, value, value_linear`) but with important memory cost. |
| 402 | 'envelope' provides gradients only for `value` and and other outputs are |
| 403 | detached. This is useful for memory saving when only the value is needed. 'last_step' provides |
| 404 | gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values. |
| 405 | 'detach' does not compute the gradients for the Sinkhorn solver. |
| 406 | |
| 407 | Returns |
| 408 | ------- |
| 409 | res : OTResult() |
| 410 | Result of the optimization problem. The information can be obtained as follows: |
| 411 | |
| 412 | - res.plan : OT plan :math:`\mathbf{T}` |
| 413 | - res.potentials : OT dual potentials |
| 414 | - res.value : Optimal value of the optimization problem |
| 415 | - res.value_linear : Linear OT loss with the optimal OT plan |
| 416 | |
| 417 | See :any:`OTResult` for more information. |