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    /// Return the operand indices for a pairwise contraction step.
303    ///
304    /// The returned indices refer to the original inputs (`0..input_count`) and
305    /// then to intermediates (`input_count..`) produced by earlier steps.
306    ///
307    /// # Examples
308    ///
309    /// ```rust
310    /// use tenferro_einsum::{ContractionTree, Subscripts};
311    ///
312    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
313    /// let tree = ContractionTree::from_pairs(
314    ///     &subs,
315    ///     &[&[2, 2], &[2, 2], &[2, 2]],
316    ///     &[(1, 2), (0, 3)],
317    /// )
318    /// .unwrap();
319    /// assert_eq!(tree.step_pair(0), Some((1, 2)));
320    /// ```
321    #[must_use]
322    pub fn step_pair(&self, step_idx: usize) -> Option<(usize, usize)> {
323        self.steps.get(step_idx).map(|step| (step.left, step.right))
324    }
325
326    /// Return the `(lhs, rhs, output)` subscripts for a pairwise step.
327    ///
328    /// The output subscripts are the intermediate labels preserved after the
329    /// contraction, or the final output labels on the last step.
330    ///
331    /// # Examples
332    ///
333    /// ```rust
334    /// use tenferro_einsum::{ContractionTree, Subscripts};
335    ///
336    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
337    /// let tree = ContractionTree::from_pairs(
338    ///     &subs,
339    ///     &[&[2, 2], &[2, 2], &[2, 2]],
340    ///     &[(1, 2), (0, 3)],
341    /// )
342    /// .unwrap();
343    /// let (lhs, rhs, out) = tree.step_subscripts(0).unwrap();
344    /// assert_eq!(lhs, &[1, 2]);
345    /// assert_eq!(rhs, &[2, 3]);
346    /// assert_eq!(out, &[1, 3]);
347    /// ```
348    #[must_use]
349    pub fn step_subscripts(&self, step_idx: usize) -> Option<(&[u32], &[u32], &[u32])> {
350        let input_count = self.subscripts.inputs.len();
351        let step = self.steps.get(step_idx)?;
352        let result_idx = input_count + step_idx;
353        let output_subs = if step_idx + 1 == self.steps.len() {
354            &self.subscripts.output
355        } else {
356            &self.operand_subs[result_idx]
357        };
358        Some((
359            &self.operand_subs[step.left],
360            &self.operand_subs[step.right],
361            output_subs,
362        ))
363    }
364
365    /// Return the precomputed lowering plan for one pairwise contraction step.
366    ///
367    /// # Examples
368    ///
369    /// ```rust
370    /// use tenferro_einsum::{ContractionTree, Subscripts};
371    ///
372    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[0, 2]);
373    /// let tree = ContractionTree::from_pairs(&subs, &[&[2, 3], &[3, 4]], &[(0, 1)]).unwrap();
374    ///
375    /// assert_eq!(tree.step_plan(0).unwrap().gemm().m(), 2);
376    /// ```
377    #[must_use]
378    pub fn step_plan(&self, step_idx: usize) -> Option<crate::lowering::PairwiseStepPlan<'_>> {
379        self.step_plans
380            .get(step_idx)
381            .map(crate::lowering::PairwiseStepPlan::new)
382    }
383
384    #[doc(hidden)]
385    #[must_use]
386    pub(crate) fn retained_bytes_for_cache_stats(&self) -> usize {
387        saturating_sum([
388            size_of::<Self>(),
389            subscripts_retained_bytes(&self.subscripts),
390            self.steps
391                .capacity()
392                .saturating_mul(size_of::<ContractionStep>()),
393            self.size_dict
394                .capacity()
395                .saturating_mul(size_of::<u32>().saturating_add(size_of::<usize>())),
396            vec_of_vec_retained_bytes(&self.operand_subs),
397            self.step_plans
398                .capacity()
399                .saturating_mul(size_of::<StepPlan>()),
400            saturating_sum(self.step_plans.iter().map(step_plan_retained_bytes)),
401        ])
402    }
403}
404
405fn subscripts_retained_bytes(subscripts: &Subscripts) -> usize {
406    saturating_sum([
407        vec_of_vec_retained_bytes(&subscripts.inputs),
408        vec_retained_bytes(&subscripts.output),
409    ])
410}
411
412fn reduce_plan_retained_bytes(plan: &ReducePlan) -> usize {
413    saturating_sum([
414        vec_retained_bytes(&plan.original_subs),
415        vec_retained_bytes(&plan.kept_subs),
416        vec_retained_bytes(&plan.out_shape),
417    ])
418}
419
420fn diag_plan_retained_bytes(plan: &DiagPlan) -> usize {
421    saturating_sum([
422        vec_retained_bytes(&plan.stages),
423        saturating_sum(plan.stages.iter().map(|stage| {
424            saturating_sum([
425                vec_retained_bytes(&stage.axis_pairs),
426                vec_retained_bytes(&stage.result_subs),
427            ])
428        })),
429        vec_retained_bytes(&plan.result_subs),
430    ])
431}
432
433fn gemm_plan_retained_bytes(plan: &GemmPlan) -> usize {
434    saturating_sum([
435        plan.reduce_a.as_ref().map_or(0, reduce_plan_retained_bytes),
436        plan.reduce_b.as_ref().map_or(0, reduce_plan_retained_bytes),
437        vec_retained_bytes(&plan.subs_a),
438        vec_retained_bytes(&plan.subs_b),
439        vec_retained_bytes(&plan.lo_modes),
440        vec_retained_bytes(&plan.ro_modes),
441        vec_retained_bytes(&plan.sum_modes),
442        vec_retained_bytes(&plan.lo_sizes),
443        vec_retained_bytes(&plan.ro_sizes),
444        vec_retained_bytes(&plan.sum_sizes),
445        vec_retained_bytes(&plan.batch_sizes),
446        vec_retained_bytes(&plan.target_a),
447        vec_retained_bytes(&plan.target_b),
448        vec_retained_bytes(&plan.c_gemm_shape),
449        vec_retained_bytes(&plan.expanded_shape),
450        vec_retained_bytes(&plan.canonical_modes),
451        vec_retained_bytes(&plan.a_gemm_shape),
452        vec_retained_bytes(&plan.b_gemm_shape),
453    ])
454}
455
456fn step_plan_retained_bytes(plan: &StepPlan) -> usize {
457    saturating_sum([
458        plan.diag_a.as_ref().map_or(0, diag_plan_retained_bytes),
459        plan.diag_b.as_ref().map_or(0, diag_plan_retained_bytes),
460        plan.strict_binary.as_ref().map_or(0, size_of_val),
461        gemm_plan_retained_bytes(&plan.gemm),
462    ])
463}
464
465fn optimize_omeco_pairs(
466    subscripts: &Subscripts,
467    size_dict: &HashMap<u32, usize>,
468    options: &ContractionOptimizerOptions,
469) -> Result<Option<Vec<(usize, usize)>>> {
470    let code = OmecoEinCode::new(subscripts.inputs.clone(), subscripts.output.clone());
471    let optimizer = options.to_treesa();
472    let Some(nested) = optimizer.optimize(&code, size_dict) else {
473        return Ok(None);
474    };
475
476    let mut next_operand = subscripts.inputs.len();
477    let mut pairs = Vec::with_capacity(subscripts.inputs.len().saturating_sub(1));
478    nested_to_pairs(&nested, &mut next_operand, &mut pairs)?;
479    Ok(Some(pairs))
480}
481
482fn nested_to_pairs(
483    nested: &NestedEinsum<u32>,
484    next_operand: &mut usize,
485    pairs: &mut Vec<(usize, usize)>,
486) -> Result<usize> {
487    match nested {
488        NestedEinsum::Leaf { tensor_index } => Ok(*tensor_index),
489        NestedEinsum::Node { args, .. } => {
490            if args.len() != 2 {
491                return Err(Error::planning(format!(
492                    "omeco returned non-binary contraction node with {} children",
493                    args.len()
494                )));
495            }
496            let left = nested_to_pairs(&args[0], next_operand, pairs)?;
497            let right = nested_to_pairs(&args[1], next_operand, pairs)?;
498            pairs.push((left, right));
499            let result_idx = *next_operand;
500            *next_operand += 1;
501            Ok(result_idx)
502        }
503    }
504}
505
506fn build_operand_label_sets(operand_subs: &[Vec<u32>]) -> Vec<HashSet<u32>> {
507    operand_subs
508        .iter()
509        .map(|subs| subs.iter().copied().collect())
510        .collect()
511}
512
513fn build_needed_label_counts(
514    output_subs: &[u32],
515    available: &[usize],
516    operand_label_sets: &[HashSet<u32>],
517) -> HashMap<u32, usize> {
518    let mut counts = HashMap::new();
519    for &label in output_subs {
520        counts.entry(label).or_insert(1);
521    }
522    for &idx in available {
523        add_labels_to_counts(&mut counts, &operand_label_sets[idx]);
524    }
525    counts
526}
527
528fn add_labels_to_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
529    for &label in labels {
530        *counts.entry(label).or_insert(0) += 1;
531    }
532}
533
534fn remove_labels_from_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
535    for &label in labels {
536        match counts.get(&label).copied() {
537            Some(1) => {
538                counts.remove(&label);
539            }
540            Some(count) => {
541                counts.insert(label, count - 1);
542            }
543            None => {}
544        }
545    }
546}
547
548fn candidate_label_is_needed(
549    label: u32,
550    left: usize,
551    right: usize,
552    operand_label_sets: &[HashSet<u32>],
553    needed_label_counts: &HashMap<u32, usize>,
554) -> bool {
555    let mut selected_count = 0;
556    if operand_label_sets[left].contains(&label) {
557        selected_count += 1;
558    }
559    if operand_label_sets[right].contains(&label) {
560        selected_count += 1;
561    }
562    needed_label_counts.get(&label).copied().unwrap_or(0) > selected_count
563}
564
565fn collect_candidate_intermediate_subs(
566    subs_left: &[u32],
567    subs_right: &[u32],
568    left: usize,
569    right: usize,
570    operand_label_sets: &[HashSet<u32>],
571    needed_label_counts: &HashMap<u32, usize>,
572    output: &mut Vec<u32>,
573) {
574    output.clear();
575    for &label in subs_left.iter().chain(subs_right.iter()) {
576        if candidate_label_is_needed(label, left, right, operand_label_sets, needed_label_counts)
577            && !output.contains(&label)
578        {
579            output.push(label);
580        }
581    }
582}
583
584#[derive(Clone, Copy)]
585struct CandidateCostContext<'a> {
586    operand_label_sets: &'a [HashSet<u32>],
587    needed_label_counts: &'a HashMap<u32, usize>,
588    size_dict: &'a HashMap<u32, usize>,
589}
590
591fn candidate_contraction_cost(
592    subs_left: &[u32],
593    subs_right: &[u32],
594    left: usize,
595    right: usize,
596    context: CandidateCostContext<'_>,
597    candidate_subs: &mut Vec<u32>,
598) -> Result<usize> {
599    collect_candidate_intermediate_subs(
600        subs_left,
601        subs_right,
602        left,
603        right,
604        context.operand_label_sets,
605        context.needed_label_counts,
606        candidate_subs,
607    );
608    let mut cost = 1usize;
609    for &label in candidate_subs.iter() {
610        let size = context.size_dict.get(&label).copied().ok_or_else(|| {
611            Error::planning(format!(
612                "unknown size for label {label} in contraction cost"
613            ))
614        })?;
615        cost = cost.saturating_mul(size);
616    }
617    Ok(cost.max(1))
618}
619
620fn optimize_self_greedy_pairs(
621    subscripts: &Subscripts,
622    size_dict: &HashMap<u32, usize>,
623) -> Result<Vec<(usize, usize)>> {
624    let input_count = subscripts.inputs.len();
625    let mut available: Vec<usize> = (0..input_count).collect();
626    let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
627    let mut operand_label_sets = build_operand_label_sets(&operand_subs);
628    let mut needed_label_counts =
629        build_needed_label_counts(&subscripts.output, &available, &operand_label_sets);
630    let mut candidate_subs = Vec::new();
631    let mut pairs: Vec<(usize, usize)> = Vec::new();
632
633    while available.len() > 1 {
634        let mut best_i = 0;
635        let mut best_j = 1;
636        let mut best_cost = usize::MAX;
637
638        for i in 0..available.len() {
639            for j in (i + 1)..available.len() {
640                let li = available[i];
641                let lj = available[j];
642                let cost = candidate_contraction_cost(
643                    &operand_subs[li],
644                    &operand_subs[lj],
645                    li,
646                    lj,
647                    CandidateCostContext {
648                        operand_label_sets: &operand_label_sets,
649                        needed_label_counts: &needed_label_counts,
650                        size_dict,
651                    },
652                    &mut candidate_subs,
653                )?;
654                if cost < best_cost {
655                    best_cost = cost;
656                    best_i = i;
657                    best_j = j;
658                }
659            }
660        }
661
662        let left = available[best_i];
663        let right = available[best_j];
664        pairs.push((left, right));
665
666        let mut new_subs = Vec::new();
667        collect_candidate_intermediate_subs(
668            &operand_subs[left],
669            &operand_subs[right],
670            left,
671            right,
672            &operand_label_sets,
673            &needed_label_counts,
674            &mut new_subs,
675        );
676        let new_idx = operand_subs.len();
677        let new_label_set: HashSet<u32> = new_subs.iter().copied().collect();
678        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[left]);
679        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[right]);
680        add_labels_to_counts(&mut needed_label_counts, &new_label_set);
681        operand_subs.push(new_subs);
682        operand_label_sets.push(new_label_set);
683        available.remove(best_j);
684        available.remove(best_i);
685        available.push(new_idx);
686    }
687
688    Ok(pairs)
689}
690
691#[cfg(test)]
692mod tests;
tenferro_einsum/planning/tree.rs

tenferro_einsum/planning/
tree.rs
