tenferro_einsum/planning/
tree.rs

1use std::collections::{HashMap, HashSet};
2use std::fmt;
3use std::mem::{size_of, size_of_val};
4
5use omeco::{
6    CodeOptimizer, EinCode as OmecoEinCode, Initializer, NestedEinsum, ScoreFunction, TreeSA,
7};
8
9use crate::cache::{saturating_sum, vec_of_vec_retained_bytes, vec_retained_bytes};
10use crate::planning::plan::{compile_step_plans, DiagPlan, GemmPlan, ReducePlan, StepPlan};
11use crate::syntax::subscripts::Subscripts;
12use crate::util::{build_size_dict, intermediate_subs};
13use crate::{Error, Result};
14
15/// A single step in the contraction sequence.
16pub(crate) struct ContractionStep {
17    pub(crate) left: usize,
18    pub(crate) right: usize,
19}
20
21/// Public options for automatic contraction-path optimization.
22///
23/// The default planner uses TreeSA with a greedy initializer and zero annealing
24/// iterations. This keeps the public API on a single optimizer family while
25/// making the default behavior effectively "greedy-only".
26#[derive(Debug, Clone)]
27pub struct ContractionOptimizerOptions {
28    /// Inverse-temperature schedule for TreeSA.
29    pub betas: Vec<f64>,
30    /// Number of independent TreeSA trials.
31    pub ntrials: usize,
32    /// Annealing iterations per temperature level.
33    pub niters: usize,
34    /// Score function used by TreeSA.
35    pub score: ScoreFunction,
36}
37
38impl Default for ContractionOptimizerOptions {
39    fn default() -> Self {
40        Self {
41            betas: Vec::new(),
42            ntrials: 1,
43            niters: 0,
44            score: ScoreFunction::default(),
45        }
46    }
47}
48
49impl ContractionOptimizerOptions {
50    fn to_treesa(&self) -> TreeSA {
51        TreeSA::new(
52            self.betas.clone(),
53            self.ntrials,
54            self.niters,
55            Initializer::Greedy,
56            self.score.clone(),
57        )
58    }
59
60    pub(crate) fn validate(&self) -> Result<()> {
61        if self.ntrials == 0 {
62            return Err(Error::planning(
63                "contraction optimizer ntrials must be at least 1",
64            ));
65        }
66        if self.betas.iter().any(|value| value.is_nan()) {
67            return Err(Error::planning(
68                "contraction optimizer betas must not contain NaN",
69            ));
70        }
71        if self.score.tc_weight.is_nan()
72            || self.score.sc_weight.is_nan()
73            || self.score.rw_weight.is_nan()
74            || self.score.sc_target.is_nan()
75        {
76            return Err(Error::planning(
77                "contraction optimizer score fields must not contain NaN",
78            ));
79        }
80        Ok(())
81    }
82}
83
84/// Contraction tree determining pairwise contraction order for N-ary einsum.
85///
86/// When contracting more than two tensors, the order in which pairwise
87/// contractions are performed significantly affects performance.
88/// `ContractionTree` encodes this order as a binary tree.
89///
90/// # Optimization
91///
92/// Use [`ContractionTree::optimize`] for automatic cost-based optimization
93/// (e.g., greedy algorithm based on tensor sizes), or
94/// [`ContractionTree::from_pairs`] for manual specification.
95pub struct ContractionTree {
96    /// Original subscripts.
97    pub(crate) subscripts: Subscripts,
98    /// Steps in the contraction (empty for single-tensor case).
99    pub(crate) steps: Vec<ContractionStep>,
100    /// Label → dimension size mapping.
101    pub(crate) size_dict: HashMap<u32, usize>,
102    /// Subscripts for each operand (0..input_count from input, then intermediates).
103    pub(crate) operand_subs: Vec<Vec<u32>>,
104    /// Pre-compiled step plans (cached to avoid recomputation per execute call).
105    pub(crate) step_plans: Vec<StepPlan>,
106}
107
108impl fmt::Debug for ContractionTree {
109    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110        f.debug_struct("ContractionTree")
111            .field("input_count", &self.subscripts.inputs.len())
112            .field("output_rank", &self.subscripts.output.len())
113            .field("steps_len", &self.steps.len())
114            .field("size_dict_len", &self.size_dict.len())
115            .field("operand_subs_len", &self.operand_subs.len())
116            .field("step_plans_len", &self.step_plans.len())
117            .finish_non_exhaustive()
118    }
119}
120
121impl ContractionTree {
122    /// Automatically compute an optimized contraction order.
123    ///
124    /// Uses a cost-based heuristic (greedy algorithm) to determine
125    /// the pairwise contraction sequence that minimizes total operation count.
126    ///
127    /// # Arguments
128    ///
129    /// * `subscripts` — Einsum subscripts for all tensors
130    /// * `shapes` — Shape of each input tensor
131    ///
132    /// # Errors
133    ///
134    /// Returns [`Error::Validation`] when subscripts and shapes have different
135    /// ranks or incompatible dimensions, or [`Error::Planning`] when no valid
136    /// contraction order can be constructed.
137    pub fn optimize(subscripts: &Subscripts, shapes: &[&[usize]]) -> Result<Self> {
138        Self::optimize_with_options(subscripts, shapes, &ContractionOptimizerOptions::default())
139    }
140
141    /// Automatically compute an optimized contraction order with explicit
142    /// planner options.
143    ///
144    /// This routes automatic planning through TreeSA using the provided
145    /// configuration. The default options correspond to a greedy-initialized
146    /// TreeSA with zero annealing iterations.
147    ///
148    /// # Errors
149    ///
150    /// Returns [`Error::Validation`] for rank, shape, or dimension mismatches,
151    /// or [`Error::Planning`] when planner options such as `ntrials` are
152    /// invalid or no contraction order can be constructed.
153    pub fn optimize_with_options(
154        subscripts: &Subscripts,
155        shapes: &[&[usize]],
156        options: &ContractionOptimizerOptions,
157    ) -> Result<Self> {
158        options.validate()?;
159        let input_count = subscripts.inputs.len();
160        if input_count <= 1 {
161            return Self::from_pairs(subscripts, shapes, &[]);
162        }
163
164        let size_dict = build_size_dict(subscripts, shapes, None)?;
165        let pairs =
166            if let Some(omeco_pairs) = optimize_omeco_pairs(subscripts, &size_dict, options)? {
167                omeco_pairs
168            } else {
169                optimize_self_greedy_pairs(subscripts, &size_dict)?
170            };
171        Self::from_pairs(subscripts, shapes, &pairs)
172    }
173
174    /// Manually build a contraction tree from a pairwise contraction sequence.
175    ///
176    /// Each pair `(i, j)` specifies which two tensors (or intermediate results)
177    /// to contract next. Intermediate results are assigned indices starting
178    /// from the number of input tensors.
179    ///
180    /// # Arguments
181    ///
182    /// * `subscripts` — Einsum subscripts for all tensors
183    /// * `shapes` — Shape of each input tensor
184    /// * `pairs` — Ordered list of pairwise contractions
185    ///
186    /// # Examples
187    ///
188    /// ```rust
189    /// use tenferro_einsum::{ContractionTree, Subscripts};
190    ///
191    /// // Three tensors: A[ij] B[jk] C[kl] -> D[il]
192    /// // Contract B and C first, then A with the result:
193    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
194    /// let shapes = [&[3, 4][..], &[4, 5], &[5, 6]];
195    /// let tree = ContractionTree::from_pairs(
196    ///     &subs,
197    ///     &shapes,
198    ///     &[(1, 2), (0, 3)],  // B*C -> T(index=3), then A*T -> D
199    /// ).unwrap();
200    /// ```
201    ///
202    /// # Errors
203    ///
204    /// Returns [`Error::Planning`] when the pair count, operand indices, or
205    /// intermediate sequence is invalid, or [`Error::Validation`] when the
206    /// supplied shapes do not match the subscripts.
207    pub fn from_pairs(
208        subscripts: &Subscripts,
209        shapes: &[&[usize]],
210        pairs: &[(usize, usize)],
211    ) -> Result<Self> {
212        let input_count = subscripts.inputs.len();
213        let required_steps = input_count.saturating_sub(1);
214        if pairs.len() != required_steps {
215            return Err(Error::planning(format!(
216                "explicit contraction path for {input_count} operands must have {required_steps} steps, got {}",
217                pairs.len()
218            )));
219        }
220        let size_dict = build_size_dict(subscripts, shapes, None)?;
221
222        let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
223        let mut live = vec![false; input_count + pairs.len()];
224        for slot in live.iter_mut().take(input_count) {
225            *slot = true;
226        }
227        let mut steps = Vec::new();
228
229        for (step_idx, &(left, right)) in pairs.iter().enumerate() {
230            let next_idx = input_count + step_idx;
231            if left == right {
232                return Err(Error::planning(format!(
233                    "pair ({left}, {right}) must reference two distinct live operands"
234                )));
235            }
236            if left >= next_idx || right >= next_idx {
237                return Err(Error::planning(format!(
238                    "pair ({left}, {right}) references non-existent operand"
239                )));
240            }
241            if !live[left] || !live[right] {
242                return Err(Error::planning(format!(
243                    "pair ({left}, {right}) references an operand or intermediate that is no longer live"
244                )));
245            }
246
247            // Labels needed by other live operands + final output
248            let mut needed: HashSet<u32> = subscripts.output.iter().copied().collect();
249            for (idx, subs) in operand_subs.iter().enumerate() {
250                if idx != left && idx != right && live[idx] {
251                    needed.extend(subs.iter().copied());
252                }
253            }
254
255            let new_subs = intermediate_subs(&operand_subs[left], &operand_subs[right], &needed);
256            operand_subs.push(new_subs);
257            live[left] = false;
258            live[right] = false;
259            live[next_idx] = true;
260            steps.push(ContractionStep { left, right });
261        }
262
263        let live_count = live.iter().filter(|&&is_live| is_live).count();
264        if live_count != 1 {
265            return Err(Error::planning(format!(
266                "explicit contraction path must leave exactly one live result, got {live_count}"
267            )));
268        }
269
270        let mut tree = Self {
271            subscripts: subscripts.clone(),
272            steps,
273            size_dict,
274            operand_subs,
275            step_plans: Vec::new(),
276        };
277        tree.step_plans = compile_step_plans(&tree)?;
278        Ok(tree)
279    }
280
281    /// Return the number of pairwise contraction steps in this tree.
282    ///
283    /// # Examples
284    ///
285    /// ```rust
286    /// use tenferro_einsum::{ContractionTree, Subscripts};
287    ///
288    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
289    /// let tree = ContractionTree::from_pairs(
290    ///     &subs,
291    ///     &[&[2, 2], &[2, 2], &[2, 2]],
292    ///     &[(1, 2), (0, 3)],
293    /// )
294    /// .unwrap();
295    /// assert_eq!(tree.step_count(), 2);
296    /// ```
297    #[must_use]
298    pub fn step_count(&self) -> usize {
299        self.steps.len()
300    }
301
302    pub(crate) fn label_size(&self, label: u32) -> Option<usize> {
303        self.size_dict.get(&label).copied()
304    }
305
306    pub(crate) fn output_shape(&self) -> Vec<usize> {
307        self.subscripts
308            .output
309            .iter()
310            .filter_map(|label| self.label_size(*label))
311            .collect()
312    }
313
314    /// Return the operand indices for a pairwise contraction step.
315    ///
316    /// The returned indices refer to the original inputs (`0..input_count`) and
317    /// then to intermediates (`input_count..`) produced by earlier steps.
318    ///
319    /// # Examples
320    ///
321    /// ```rust
322    /// use tenferro_einsum::{ContractionTree, Subscripts};
323    ///
324    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
325    /// let tree = ContractionTree::from_pairs(
326    ///     &subs,
327    ///     &[&[2, 2], &[2, 2], &[2, 2]],
328    ///     &[(1, 2), (0, 3)],
329    /// )
330    /// .unwrap();
331    /// assert_eq!(tree.step_pair(0), Some((1, 2)));
332    /// ```
333    #[must_use]
334    pub fn step_pair(&self, step_idx: usize) -> Option<(usize, usize)> {
335        self.steps.get(step_idx).map(|step| (step.left, step.right))
336    }
337
338    /// Return the `(lhs, rhs, output)` subscripts for a pairwise step.
339    ///
340    /// The output subscripts are the intermediate labels preserved after the
341    /// contraction, or the final output labels on the last step.
342    ///
343    /// # Examples
344    ///
345    /// ```rust
346    /// use tenferro_einsum::{ContractionTree, Subscripts};
347    ///
348    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
349    /// let tree = ContractionTree::from_pairs(
350    ///     &subs,
351    ///     &[&[2, 2], &[2, 2], &[2, 2]],
352    ///     &[(1, 2), (0, 3)],
353    /// )
354    /// .unwrap();
355    /// let (lhs, rhs, out) = tree.step_subscripts(0).unwrap();
356    /// assert_eq!(lhs, &[1, 2]);
357    /// assert_eq!(rhs, &[2, 3]);
358    /// assert_eq!(out, &[1, 3]);
359    /// ```
360    #[must_use]
361    pub fn step_subscripts(&self, step_idx: usize) -> Option<(&[u32], &[u32], &[u32])> {
362        let input_count = self.subscripts.inputs.len();
363        let step = self.steps.get(step_idx)?;
364        let result_idx = input_count + step_idx;
365        let output_subs = if step_idx + 1 == self.steps.len() {
366            &self.subscripts.output
367        } else {
368            &self.operand_subs[result_idx]
369        };
370        Some((
371            &self.operand_subs[step.left],
372            &self.operand_subs[step.right],
373            output_subs,
374        ))
375    }
376
377    /// Return the precomputed lowering plan for one pairwise contraction step.
378    ///
379    /// # Examples
380    ///
381    /// ```rust
382    /// use tenferro_einsum::{ContractionTree, Subscripts};
383    ///
384    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[0, 2]);
385    /// let tree = ContractionTree::from_pairs(&subs, &[&[2, 3], &[3, 4]], &[(0, 1)]).unwrap();
386    ///
387    /// assert_eq!(tree.step_plan(0).unwrap().gemm().m(), 2);
388    /// ```
389    #[must_use]
390    pub fn step_plan(&self, step_idx: usize) -> Option<crate::lowering::PairwiseStepPlan<'_>> {
391        self.step_plans
392            .get(step_idx)
393            .map(crate::lowering::PairwiseStepPlan::new)
394    }
395
396    #[doc(hidden)]
397    #[must_use]
398    pub(crate) fn retained_bytes_for_cache_stats(&self) -> usize {
399        saturating_sum([
400            size_of::<Self>(),
401            subscripts_retained_bytes(&self.subscripts),
402            self.steps
403                .capacity()
404                .saturating_mul(size_of::<ContractionStep>()),
405            self.size_dict
406                .capacity()
407                .saturating_mul(size_of::<u32>().saturating_add(size_of::<usize>())),
408            vec_of_vec_retained_bytes(&self.operand_subs),
409            self.step_plans
410                .capacity()
411                .saturating_mul(size_of::<StepPlan>()),
412            saturating_sum(self.step_plans.iter().map(step_plan_retained_bytes)),
413        ])
414    }
415}
416
417fn subscripts_retained_bytes(subscripts: &Subscripts) -> usize {
418    saturating_sum([
419        vec_of_vec_retained_bytes(&subscripts.inputs),
420        vec_retained_bytes(&subscripts.output),
421    ])
422}
423
424fn reduce_plan_retained_bytes(plan: &ReducePlan) -> usize {
425    saturating_sum([
426        vec_retained_bytes(&plan.original_subs),
427        vec_retained_bytes(&plan.kept_subs),
428        vec_retained_bytes(&plan.out_shape),
429    ])
430}
431
432fn diag_plan_retained_bytes(plan: &DiagPlan) -> usize {
433    saturating_sum([
434        vec_retained_bytes(&plan.stages),
435        saturating_sum(plan.stages.iter().map(|stage| {
436            saturating_sum([
437                vec_retained_bytes(&stage.axis_pairs),
438                vec_retained_bytes(&stage.result_subs),
439            ])
440        })),
441        vec_retained_bytes(&plan.result_subs),
442    ])
443}
444
445fn gemm_plan_retained_bytes(plan: &GemmPlan) -> usize {
446    saturating_sum([
447        plan.reduce_a.as_ref().map_or(0, reduce_plan_retained_bytes),
448        plan.reduce_b.as_ref().map_or(0, reduce_plan_retained_bytes),
449        vec_retained_bytes(&plan.subs_a),
450        vec_retained_bytes(&plan.subs_b),
451        vec_retained_bytes(&plan.lo_modes),
452        vec_retained_bytes(&plan.ro_modes),
453        vec_retained_bytes(&plan.sum_modes),
454        vec_retained_bytes(&plan.lo_sizes),
455        vec_retained_bytes(&plan.ro_sizes),
456        vec_retained_bytes(&plan.sum_sizes),
457        vec_retained_bytes(&plan.batch_sizes),
458        vec_retained_bytes(&plan.target_a),
459        vec_retained_bytes(&plan.target_b),
460        vec_retained_bytes(&plan.c_gemm_shape),
461        vec_retained_bytes(&plan.expanded_shape),
462        vec_retained_bytes(&plan.canonical_modes),
463        vec_retained_bytes(&plan.a_gemm_shape),
464        vec_retained_bytes(&plan.b_gemm_shape),
465    ])
466}
467
468fn step_plan_retained_bytes(plan: &StepPlan) -> usize {
469    saturating_sum([
470        plan.diag_a.as_ref().map_or(0, diag_plan_retained_bytes),
471        plan.diag_b.as_ref().map_or(0, diag_plan_retained_bytes),
472        plan.strict_binary.as_ref().map_or(0, size_of_val),
473        gemm_plan_retained_bytes(&plan.gemm),
474    ])
475}
476
477fn optimize_omeco_pairs(
478    subscripts: &Subscripts,
479    size_dict: &HashMap<u32, usize>,
480    options: &ContractionOptimizerOptions,
481) -> Result<Option<Vec<(usize, usize)>>> {
482    let code = OmecoEinCode::new(subscripts.inputs.clone(), subscripts.output.clone());
483    let optimizer = options.to_treesa();
484    let Some(nested) = optimizer.optimize(&code, size_dict) else {
485        return Ok(None);
486    };
487
488    let mut next_operand = subscripts.inputs.len();
489    let mut pairs = Vec::with_capacity(subscripts.inputs.len().saturating_sub(1));
490    nested_to_pairs(&nested, &mut next_operand, &mut pairs)?;
491    Ok(Some(pairs))
492}
493
494fn nested_to_pairs(
495    nested: &NestedEinsum<u32>,
496    next_operand: &mut usize,
497    pairs: &mut Vec<(usize, usize)>,
498) -> Result<usize> {
499    match nested {
500        NestedEinsum::Leaf { tensor_index } => Ok(*tensor_index),
501        NestedEinsum::Node { args, .. } => {
502            if args.len() != 2 {
503                return Err(Error::planning(format!(
504                    "omeco returned non-binary contraction node with {} children",
505                    args.len()
506                )));
507            }
508            let left = nested_to_pairs(&args[0], next_operand, pairs)?;
509            let right = nested_to_pairs(&args[1], next_operand, pairs)?;
510            pairs.push((left, right));
511            let result_idx = *next_operand;
512            *next_operand += 1;
513            Ok(result_idx)
514        }
515    }
516}
517
518fn build_operand_label_sets(operand_subs: &[Vec<u32>]) -> Vec<HashSet<u32>> {
519    operand_subs
520        .iter()
521        .map(|subs| subs.iter().copied().collect())
522        .collect()
523}
524
525fn build_needed_label_counts(
526    output_subs: &[u32],
527    available: &[usize],
528    operand_label_sets: &[HashSet<u32>],
529) -> HashMap<u32, usize> {
530    let mut counts = HashMap::new();
531    for &label in output_subs {
532        counts.entry(label).or_insert(1);
533    }
534    for &idx in available {
535        add_labels_to_counts(&mut counts, &operand_label_sets[idx]);
536    }
537    counts
538}
539
540fn add_labels_to_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
541    for &label in labels {
542        *counts.entry(label).or_insert(0) += 1;
543    }
544}
545
546fn remove_labels_from_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
547    for &label in labels {
548        match counts.get(&label).copied() {
549            Some(1) => {
550                counts.remove(&label);
551            }
552            Some(count) => {
553                counts.insert(label, count - 1);
554            }
555            None => {}
556        }
557    }
558}
559
560fn candidate_label_is_needed(
561    label: u32,
562    left: usize,
563    right: usize,
564    operand_label_sets: &[HashSet<u32>],
565    needed_label_counts: &HashMap<u32, usize>,
566) -> bool {
567    let mut selected_count = 0;
568    if operand_label_sets[left].contains(&label) {
569        selected_count += 1;
570    }
571    if operand_label_sets[right].contains(&label) {
572        selected_count += 1;
573    }
574    needed_label_counts.get(&label).copied().unwrap_or(0) > selected_count
575}
576
577fn collect_candidate_intermediate_subs(
578    subs_left: &[u32],
579    subs_right: &[u32],
580    left: usize,
581    right: usize,
582    operand_label_sets: &[HashSet<u32>],
583    needed_label_counts: &HashMap<u32, usize>,
584    output: &mut Vec<u32>,
585) {
586    output.clear();
587    for &label in subs_left.iter().chain(subs_right.iter()) {
588        if candidate_label_is_needed(label, left, right, operand_label_sets, needed_label_counts)
589            && !output.contains(&label)
590        {
591            output.push(label);
592        }
593    }
594}
595
596#[derive(Clone, Copy)]
597struct CandidateCostContext<'a> {
598    operand_label_sets: &'a [HashSet<u32>],
599    needed_label_counts: &'a HashMap<u32, usize>,
600    size_dict: &'a HashMap<u32, usize>,
601}
602
603fn candidate_contraction_cost(
604    subs_left: &[u32],
605    subs_right: &[u32],
606    left: usize,
607    right: usize,
608    context: CandidateCostContext<'_>,
609    candidate_subs: &mut Vec<u32>,
610) -> Result<usize> {
611    collect_candidate_intermediate_subs(
612        subs_left,
613        subs_right,
614        left,
615        right,
616        context.operand_label_sets,
617        context.needed_label_counts,
618        candidate_subs,
619    );
620    let mut cost = 1usize;
621    for &label in candidate_subs.iter() {
622        let size = context.size_dict.get(&label).copied().ok_or_else(|| {
623            Error::planning(format!(
624                "unknown size for label {label} in contraction cost"
625            ))
626        })?;
627        cost = cost.saturating_mul(size);
628    }
629    Ok(cost.max(1))
630}
631
632fn optimize_self_greedy_pairs(
633    subscripts: &Subscripts,
634    size_dict: &HashMap<u32, usize>,
635) -> Result<Vec<(usize, usize)>> {
636    let input_count = subscripts.inputs.len();
637    let mut available: Vec<usize> = (0..input_count).collect();
638    let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
639    let mut operand_label_sets = build_operand_label_sets(&operand_subs);
640    let mut needed_label_counts =
641        build_needed_label_counts(&subscripts.output, &available, &operand_label_sets);
642    let mut candidate_subs = Vec::new();
643    let mut pairs: Vec<(usize, usize)> = Vec::new();
644
645    while available.len() > 1 {
646        let mut best_i = 0;
647        let mut best_j = 1;
648        let mut best_cost = usize::MAX;
649
650        for i in 0..available.len() {
651            for j in (i + 1)..available.len() {
652                let li = available[i];
653                let lj = available[j];
654                let cost = candidate_contraction_cost(
655                    &operand_subs[li],
656                    &operand_subs[lj],
657                    li,
658                    lj,
659                    CandidateCostContext {
660                        operand_label_sets: &operand_label_sets,
661                        needed_label_counts: &needed_label_counts,
662                        size_dict,
663                    },
664                    &mut candidate_subs,
665                )?;
666                if cost < best_cost {
667                    best_cost = cost;
668                    best_i = i;
669                    best_j = j;
670                }
671            }
672        }
673
674        let left = available[best_i];
675        let right = available[best_j];
676        pairs.push((left, right));
677
678        let mut new_subs = Vec::new();
679        collect_candidate_intermediate_subs(
680            &operand_subs[left],
681            &operand_subs[right],
682            left,
683            right,
684            &operand_label_sets,
685            &needed_label_counts,
686            &mut new_subs,
687        );
688        let new_idx = operand_subs.len();
689        let new_label_set: HashSet<u32> = new_subs.iter().copied().collect();
690        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[left]);
691        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[right]);
692        add_labels_to_counts(&mut needed_label_counts, &new_label_set);
693        operand_subs.push(new_subs);
694        operand_label_sets.push(new_label_set);
695        available.remove(best_j);
696        available.remove(best_i);
697        available.push(new_idx);
698    }
699
700    Ok(pairs)
701}
702
703#[cfg(test)]
704mod tests;
tenferro_einsum/planning/tree.rs

tenferro_einsum/planning/
tree.rs
