(self, root=0, layout=None)
| 800 | return recv_buff |
| 801 | |
| 802 | def gather_data(self, root=0, layout=None): |
| 803 | # Change layout |
| 804 | if layout is not None: |
| 805 | self.change_layout(layout) |
| 806 | # Shortcut for serial execution |
| 807 | if self.dist.comm.size == 1: |
| 808 | return self.data.copy() |
| 809 | # TODO: Shortcut this for constant fields |
| 810 | # Gather data |
| 811 | # Should be optimized via Gatherv eventually |
| 812 | pieces = self.dist.comm.gather(self.data, root=root) |
| 813 | # Assemble on root node |
| 814 | if self.dist.comm.rank == root: |
| 815 | ext_mesh = self.layout.ext_mesh |
| 816 | combined = np.zeros(prod(ext_mesh), dtype=object) |
| 817 | combined[:] = pieces |
| 818 | return np.block(combined.reshape(ext_mesh).tolist()) |
| 819 | |
| 820 | def allreduce_data_norm(self, layout=None, order=2): |
| 821 | # Change layout |
no test coverage detected