Compute the cost (removed size + flops) and resultant indices for performing the contraction specified by ``positions``. Parameters ---------- positions : tuple of int The locations of the proposed tensors to contract. input_sets : list of sets The indices found
(positions, input_sets, output_set, idx_dict, memory_limit, path_cost, naive_cost)
| 213 | return path |
| 214 | |
| 215 | def _parse_possible_contraction(positions, input_sets, output_set, idx_dict, memory_limit, path_cost, naive_cost): |
| 216 | """Compute the cost (removed size + flops) and resultant indices for |
| 217 | performing the contraction specified by ``positions``. |
| 218 | |
| 219 | Parameters |
| 220 | ---------- |
| 221 | positions : tuple of int |
| 222 | The locations of the proposed tensors to contract. |
| 223 | input_sets : list of sets |
| 224 | The indices found on each tensors. |
| 225 | output_set : set |
| 226 | The output indices of the expression. |
| 227 | idx_dict : dict |
| 228 | Mapping of each index to its size. |
| 229 | memory_limit : int |
| 230 | The total allowed size for an intermediary tensor. |
| 231 | path_cost : int |
| 232 | The contraction cost so far. |
| 233 | naive_cost : int |
| 234 | The cost of the unoptimized expression. |
| 235 | |
| 236 | Returns |
| 237 | ------- |
| 238 | cost : (int, int) |
| 239 | A tuple containing the size of any indices removed, and the flop cost. |
| 240 | positions : tuple of int |
| 241 | The locations of the proposed tensors to contract. |
| 242 | new_input_sets : list of sets |
| 243 | The resulting new list of indices if this proposed contraction is performed. |
| 244 | |
| 245 | """ |
| 246 | |
| 247 | # Find the contraction |
| 248 | contract = _find_contraction(positions, input_sets, output_set) |
| 249 | idx_result, new_input_sets, idx_removed, idx_contract = contract |
| 250 | |
| 251 | # Sieve the results based on memory_limit |
| 252 | new_size = _compute_size_by_dict(idx_result, idx_dict) |
| 253 | if new_size > memory_limit: |
| 254 | return None |
| 255 | |
| 256 | # Build sort tuple |
| 257 | old_sizes = (_compute_size_by_dict(input_sets[p], idx_dict) for p in positions) |
| 258 | removed_size = sum(old_sizes) - new_size |
| 259 | |
| 260 | # NB: removed_size used to be just the size of any removed indices i.e.: |
| 261 | # helpers.compute_size_by_dict(idx_removed, idx_dict) |
| 262 | cost = _flop_count(idx_contract, idx_removed, len(positions), idx_dict) |
| 263 | sort = (-removed_size, cost) |
| 264 | |
| 265 | # Sieve based on total cost as well |
| 266 | if (path_cost + cost) > naive_cost: |
| 267 | return None |
| 268 | |
| 269 | # Add contraction to possible choices |
| 270 | return [sort, positions, new_input_sets] |
| 271 | |
| 272 |
no test coverage detected