Reduce a LazyTensor along an axis with function fun using batches. When axis=None, reduce the LazyTensor to a scalar as a sum of fun over batches taken along dim. .. warning:: This function works for tensor of any order but the reduction can be done only along the first
(a, func, axis=None, nx=None, batch_size=100)
| 634 | |
| 635 | |
| 636 | def reduce_lazytensor(a, func, axis=None, nx=None, batch_size=100): |
| 637 | """Reduce a LazyTensor along an axis with function fun using batches. |
| 638 | |
| 639 | When axis=None, reduce the LazyTensor to a scalar as a sum of fun over |
| 640 | batches taken along dim. |
| 641 | |
| 642 | .. warning:: |
| 643 | This function works for tensor of any order but the reduction can be done |
| 644 | only along the first two axis (or global). Also, in order to work, it requires that the slice of size `batch_size` along the axis to reduce (or axis 0 if `axis=None`) is can be computed and fits in memory. |
| 645 | |
| 646 | |
| 647 | Parameters |
| 648 | ---------- |
| 649 | a : LazyTensor |
| 650 | LazyTensor to reduce |
| 651 | func : callable |
| 652 | Function to apply to the LazyTensor |
| 653 | axis : int, optional |
| 654 | Axis along which to reduce the LazyTensor. If None, reduce the |
| 655 | LazyTensor to a scalar as a sum of fun over batches taken along axis 0. |
| 656 | If 0 or 1 reduce the LazyTensor to a vector/matrix as a sum of fun over |
| 657 | batches taken along axis. |
| 658 | nx : Backend, optional |
| 659 | Backend to use for the reduction |
| 660 | batch_size : int, optional |
| 661 | Size of the batches to use for the reduction (default=100) |
| 662 | |
| 663 | Returns |
| 664 | ------- |
| 665 | res : array-like |
| 666 | Result of the reduction |
| 667 | |
| 668 | """ |
| 669 | |
| 670 | if nx is None: |
| 671 | nx = get_backend(a[0:1]) |
| 672 | |
| 673 | if axis is None: |
| 674 | res = 0.0 |
| 675 | for i in range(0, a.shape[0], batch_size): |
| 676 | res += func(a[i : i + batch_size]) |
| 677 | return res |
| 678 | elif axis == 0: |
| 679 | res = nx.zeros(a.shape[1:], type_as=a[0]) |
| 680 | if nx.__name__ in ["jax", "tf"]: |
| 681 | lst = [] |
| 682 | for j in range(0, a.shape[1], batch_size): |
| 683 | lst.append(func(a[:, j : j + batch_size], 0)) |
| 684 | return nx.concatenate(lst, axis=0) |
| 685 | else: |
| 686 | for j in range(0, a.shape[1], batch_size): |
| 687 | res[j : j + batch_size] = func(a[:, j : j + batch_size], axis=0) |
| 688 | return res |
| 689 | elif axis == 1: |
| 690 | if len(a.shape) == 2: |
| 691 | shape = a.shape[0] |
| 692 | else: |
| 693 | shape = (a.shape[0], *a.shape[2:]) |
no test coverage detected