Parse broadcast and core dimensions for vectorize with a signature. Arguments --------- args : Tuple[ndarray, ...] Tuple of input arguments to examine. input_core_dims : List[Tuple[str, ...]] List of core dimensions corresponding to each input. Returns
(args, input_core_dims)
| 2217 | |
| 2218 | |
| 2219 | def _parse_input_dimensions(args, input_core_dims): |
| 2220 | """ |
| 2221 | Parse broadcast and core dimensions for vectorize with a signature. |
| 2222 | |
| 2223 | Arguments |
| 2224 | --------- |
| 2225 | args : Tuple[ndarray, ...] |
| 2226 | Tuple of input arguments to examine. |
| 2227 | input_core_dims : List[Tuple[str, ...]] |
| 2228 | List of core dimensions corresponding to each input. |
| 2229 | |
| 2230 | Returns |
| 2231 | ------- |
| 2232 | broadcast_shape : Tuple[int, ...] |
| 2233 | Common shape to broadcast all non-core dimensions to. |
| 2234 | dim_sizes : Dict[str, int] |
| 2235 | Common sizes for named core dimensions. |
| 2236 | """ |
| 2237 | broadcast_args = [] |
| 2238 | dim_sizes = {} |
| 2239 | for arg, core_dims in zip(args, input_core_dims): |
| 2240 | _update_dim_sizes(dim_sizes, arg, core_dims) |
| 2241 | ndim = arg.ndim - len(core_dims) |
| 2242 | dummy_array = np.lib.stride_tricks.as_strided(0, arg.shape[:ndim]) |
| 2243 | broadcast_args.append(dummy_array) |
| 2244 | broadcast_shape = np.lib._stride_tricks_impl._broadcast_shape( |
| 2245 | *broadcast_args |
| 2246 | ) |
| 2247 | return broadcast_shape, dim_sizes |
| 2248 | |
| 2249 | |
| 2250 | def _calculate_shapes(broadcast_shape, dim_sizes, list_of_core_dims): |
no test coverage detected
searching dependent graphs…