Apply a function, and enforce the output.ndim to match expected_ndim Ensures the output has the expected dimensionality.
(*args, **kwargs)
| 1024 | |
| 1025 | |
| 1026 | def apply_and_enforce(*args, **kwargs): |
| 1027 | """Apply a function, and enforce the output.ndim to match expected_ndim |
| 1028 | |
| 1029 | Ensures the output has the expected dimensionality.""" |
| 1030 | func = kwargs.pop("_func") |
| 1031 | expected_ndim = kwargs.pop("expected_ndim") |
| 1032 | out = func(*args, **kwargs) |
| 1033 | if getattr(out, "ndim", 0) != expected_ndim: |
| 1034 | out_ndim = getattr(out, "ndim", 0) |
| 1035 | raise ValueError( |
| 1036 | f"Dimension mismatch: expected output of {func} " |
| 1037 | f"to have dims = {expected_ndim}. Got {out_ndim} instead." |
| 1038 | ) |
| 1039 | return out |
| 1040 | |
| 1041 | |
| 1042 | def broadcast_chunks(*chunkss): |