tract_core/ops/array/

broadcast.rs

use tract_data::itertools::izip;

use crate::broadcast::multi_broadcast;
use crate::internal::*;
use crate::ops::binary::TypedBinOp;

#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
pub struct MultiBroadcastTo {
    pub shape: ShapeFact,
}

impl Op for MultiBroadcastTo {
    fn name(&self) -> StaticName {
        "MultiBroadcastTo".into()
    }

    op_as_typed_op!();
}

impl EvalOp for MultiBroadcastTo {
    op_out_of_plan!();

    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
        let shape = self.shape.eval_to_usize(ctx.symbols)?;
        Ok(tvec!(inputs[0].broadcast_to_shape(&shape)?.into_tvalue()))
    }
}

impl TypedOp for MultiBroadcastTo {
    fn axes_mapping(
        &self,
        inputs: &[&TypedFact],
        outputs: &[&TypedFact],
    ) -> TractResult<AxesMapping> {
        // ONNX-style broadcasting right-aligns input over output, so when
        // output_rank > input_rank the leading output axes are pure
        // broadcast axes with no input correspondence. natural_for_rank's
        // square shape would skip them and trip the optimizer's axes-mapping
        // check (caught under paranoid_assertions).
        let in_rank = inputs[0].rank();
        let out_rank = outputs[0].rank();
        let leading = out_rank.saturating_sub(in_rank);
        let mut axes = tvec!();
        let mut alphabet = 'a'..;
        for o in 0..leading {
            axes.push(
                Axis::new(alphabet.next().unwrap(), inputs.len(), outputs.len()).output(0, o),
            );
        }
        for i in 0..in_rank.min(out_rank) {
            axes.push(
                Axis::new(alphabet.next().unwrap(), inputs.len(), outputs.len())
                    .input(0, i)
                    .output(0, leading + i),
            );
        }
        AxesMapping::new(inputs.len(), outputs.len(), axes)
    }

    fn change_axes(
        &self,
        model: &TypedModel,
        node: &TypedNode,
        _io: InOut,
        change: &AxisOp,
    ) -> TractResult<Option<AxisChangeConsequence>> {
        // The output takes the change: the broadcast absorbs it in the target
        // shape. The input takes it only when every touched axis has the same
        // dim in input and target, which is a real correspondence only at
        // equal ranks (right-aligned broadcasting pairs by index then).
        let input_shape = &model.outlet_fact(node.inputs[0])?.shape;
        let canonical = change.canonical();
        let touched: TVec<usize> = match canonical.as_ref() {
            AxisOp::Add(ix) | AxisOp::Rm(ix) => tvec![*ix],
            AxisOp::Move(from, to) => {
                rule_if!(input_shape.rank() == self.shape.rank());
                tvec![*from, *to]
            }
            _ => return Ok(None),
        };
        let passthrough = touched.iter().all(|&ix| {
            ix < input_shape.rank() && ix < self.shape.rank() && input_shape[ix] == self.shape[ix]
        });
        // A Move must propagate, never absorb: transpose-of-broadcast is not
        // broadcast-into-transposed-target ([1,7] into a transposed [7,7]
        // target yields the transposed tensor).
        if matches!(canonical.as_ref(), AxisOp::Move(..)) && !passthrough {
            return Ok(None);
        }

        let mut shape = self.shape.clone();
        if change.change_shape(&mut shape, false).is_ok() {
            let mut wire_changes: TVec<(InOut, AxisOp)> = tvec![(InOut::Out(0), change.clone())];
            if passthrough {
                wire_changes.push((InOut::In(0), change.clone()));
            } else {
                // Absorbed changes must not re-pair the input's axes. The
                // input binds to the right-aligned window [offset, offset+r)
                // of the target; a change inside the window shifts input axes
                // onto neighbouring target axes (a silent permutation when
                // adjacent dims are equal, e.g. [5,1] into [5,5,1] with Rm(2)
                // absorbed), and a change of the dim-1 axes can leave the
                // input un-broadcastable into the new target ([5,1] into
                // [9,5,1] must not become [9,5]).
                let offset = self.shape.rank() as isize - input_shape.rank() as isize;
                let k = touched[0] as isize - offset;
                // the window is right-anchored: any rank change moves it, so
                // only a change at or before its start (k <= 0 for Add, which
                // shifts every partner uniformly, k < 0 for Rm, which must not
                // eat a partner) keeps the pairing. Otherwise the displaced
                // input axes must all be 1.
                let displaced_all_ones =
                    (0..k.max(0) as usize).all(|i| input_shape[i] == 1.to_dim());
                let window_safe = match canonical.as_ref() {
                    AxisOp::Add(_) => k <= 0,
                    _ => k < 0,
                } || displaced_all_ones;
                let broadcastable = multi_broadcast(&[input_shape, &shape])
                    .is_ok_and(|b| b.as_slice() == shape.as_ref());
                if !(window_safe && broadcastable) {
                    return Ok(None);
                }
            }
            return Ok(Some(AxisChangeConsequence {
                wire_changes,
                substitute_op: Some(Box::new(MultiBroadcastTo { shape })),
            }));
        }
        Ok(None)
    }

    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
        ensure!(inputs.len() == 1);
        let mut fact = inputs[0].datum_type.fact(self.shape.clone());
        fact.uniform.clone_from(&inputs[0].uniform);
        fact.uniform_tdim = inputs[0].uniform_tdim.clone();
        Ok(tvec!(fact))
    }

    fn input_roi(
        &self,
        model: &TypedModel,
        node: &TypedNode,
    ) -> TractResult<Option<TVec<Option<TDim>>>> {
        crate::optim::propagate_roi::bubble_roi(model, node)
    }

    fn set_symbols(
        &self,
        _source: &TypedModel,
        node: &TypedNode,
        target: &mut TypedModel,
        mapping: &HashMap<OutletId, OutletId>,
        subs: &HashMap<Symbol, TDim>,
    ) -> TractResult<TVec<OutletId>> {
        let input = mapping[&node.inputs[0]];
        let shape: TVec<_> =
            self.shape.iter().map(|d| d.substitute_all(subs)).collect::<TractResult<_>>()?;
        let op = Self { shape: shape.into() };
        target.wire_node(&node.name, op, &[input])
    }

    fn declutter(
        &self,
        model: &TypedModel,
        node: &TypedNode,
    ) -> TractResult<Option<TypedModelPatch>> {
        let input_fact = model.outlet_fact(node.inputs[0])?;
        if input_fact.shape == self.shape {
            return TypedModelPatch::shunt_one_op(model, node);
        }
        // Swap with an AxisOp successor: `Broadcast(x, S) → AxisOp` becomes
        // `AxisOp(x) → Broadcast(σ(S))` whenever the AxisOp transforms every
        // axis the broadcast actually expanded.  Fires per-successor, so this
        // works under fan-out (the original broadcast stays in place for
        // siblings; only the matched AxisOp branch is rerouted).
        for succ in &*node.outputs[0].successors {
            let succ = model.node(succ.node);
            let Some(op) = succ.op_as::<AxisOp>() else { continue };
            // The AxisOp's indices refer to the broadcast output; they are only
            // meaningful on the input if the broadcast did not add leading axes.
            if input_fact.rank() != self.shape.rank() {
                continue;
            }
            let mut shape = self.shape.clone();
            if izip!(0.., &*input_fact.shape, &*self.shape)
                .filter(|(_, l, r)| l != r)
                .all(|(axis, _, _)| op.transform_axis(axis).is_some())
                && op.change_shape(&mut shape, false).is_ok()
            {
                let mut patch = TypedModelPatch::default();
                let mut wire = patch.tap_model(model, node.inputs[0])?;
                wire = patch.wire_node(&succ.name, op.clone(), &[wire])?[0];
                wire = patch.wire_node(&node.name, MultiBroadcastTo { shape }, &[wire])?[0];
                patch.shunt_outside(model, succ.id.into(), wire)?;
                return Ok(Some(patch));
            }
        }
        if let [succ] = &*node.outputs[0].successors {
            let succ = model.node(succ.node);
            if succ.op_is::<TypedBinOp>() {
                let our_slot = node.outputs[0].successors[0].slot;
                let other_slot = 1 - our_slot;
                let other_operand = succ.inputs[other_slot];
                let other_fact = model.outlet_fact(other_operand)?;
                let output_fact = model.outlet_fact(succ.id.into())?;
                if input_fact.rank() == other_fact.rank()
                    && multi_broadcast(&[&input_fact.shape, &other_fact.shape])
                        .is_ok_and(|s| *s == *output_fact.shape)
                {
                    let mut operands = tvec!(node.inputs[0], other_operand);
                    if our_slot == 1 {
                        operands.swap(0, 1);
                    }
                    return TypedModelPatch::rewire(
                        model,
                        &operands,
                        &[succ.id.into()],
                        &|p, inputs| p.wire_node(&succ.name, succ.op.clone(), inputs),
                    )
                    .map(Some);
                }
            }
        }
        Ok(None)
    }

    as_op!();
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ops::change_axes::AxisOp;
    use crate::ops::change_axes::{AxisChange, InOut};
    use crate::ops::logic::And;
    use proptest::prelude::*;

    /// `Broadcast → Move` with the broadcast feeding a SINGLE successor.
    /// Pre-existing path: the swap rewrite kicks in.
    #[test]
    fn broadcast_move_single_successor_swaps() -> TractResult<()> {
        let mut model = TypedModel::default();
        let t = model.symbols.sym("T");
        let pad = model.add_source("pad", bool::fact(&[t.to_dim()]))?;
        let unsq = model.wire_node("unsq", AxisOp::Add(0), &[pad])?[0];
        let bcast = model.wire_node(
            "bcast",
            MultiBroadcastTo { shape: ShapeFact::from_dims([t.to_dim(), t.to_dim()]) },
            &[unsq],
        )?[0];
        let mv = model.wire_node("move", AxisOp::Move(0, 1), &[bcast])?[0];
        model.select_output_outlets(&[mv])?;

        let model = model.into_decluttered()?;

        let move_count = model
            .nodes()
            .iter()
            .filter(|n| matches!(n.op_as::<AxisOp>(), Some(AxisOp::Move(0, 1))))
            .count();
        assert_eq!(move_count, 0, "Move should have been pushed through Broadcast and absorbed");
        Ok(())
    }

    /// `Broadcast → {Move, And-direct}` — the encoder-style pad-mask outer-AND
    /// pattern.  Pre-fix: declutter bailed because broadcast had > 1 successor;
    /// the Move stayed.  Post-fix: the Move-branch gets its own swapped
    /// chain, the direct-AND branch still consumes the original broadcast.
    #[test]
    fn broadcast_move_fanout_pushes_through_one_branch() -> TractResult<()> {
        let mut model = TypedModel::default();
        let t = model.symbols.sym("T");
        let pad = model.add_source("pad", bool::fact(&[t.to_dim()]))?;
        let unsq = model.wire_node("unsq", AxisOp::Add(0), &[pad])?[0];
        let bcast = model.wire_node(
            "bcast",
            MultiBroadcastTo { shape: ShapeFact::from_dims([t.to_dim(), t.to_dim()]) },
            &[unsq],
        )?[0];
        let mv = model.wire_node("move", AxisOp::Move(0, 1), &[bcast])?[0];
        let and = model.wire_node("and", TypedBinOp(Box::new(And), None), &[bcast, mv])?[0];
        model.select_output_outlets(&[and])?;

        let model = model.into_decluttered()?;

        // Expected: fan-out swap-through fires on the Move branch, then the
        // existing Broadcast→TypedBinOp rule fires on each (now single-
        // successor) broadcast, eliminating both — the AND ends up
        // broadcasting [1, T] and [T, 1] implicitly.
        let bcast_count = model.nodes().iter().filter(|n| n.op_is::<MultiBroadcastTo>()).count();
        assert_eq!(
            bcast_count, 0,
            "Both broadcasts should be subsumed into AND's implicit broadcasting"
        );

        let and_node =
            model.nodes().iter().find(|n| n.op_is::<TypedBinOp>()).expect("AND should survive");
        assert_eq!(and_node.inputs.len(), 2);
        let and_input_shapes: Vec<_> = and_node
            .inputs
            .iter()
            .map(|i| model.outlet_fact(*i).unwrap().shape.to_tvec())
            .collect();
        let expected_a = tvec![1.to_dim(), t.to_dim()];
        let expected_b = tvec![t.to_dim(), 1.to_dim()];
        let (a, b) = (&and_input_shapes[0], &and_input_shapes[1]);
        assert!(
            (a == &expected_a && b == &expected_b) || (a == &expected_b && b == &expected_a),
            "AND should receive [1, T] and [T, 1]; got {a:?} and {b:?}"
        );
        Ok(())
    }

    /// `Broadcast → AxisOp` where the broadcast adds a leading axis (input
    /// rank < output rank).  The AxisOp's indices refer to the output shape
    /// and are meaningless on the input; the swap must not fire.  Pre-fix,
    /// the guard izip truncated to the shorter rank and wiring the AxisOp
    /// onto the input panicked in AxisOp::change_shape.
    #[test]
    fn broadcast_adding_leading_axis_does_not_swap_with_axis_op() -> TractResult<()> {
        let mut model = TypedModel::default();
        let src = model.add_source("src", f32::fact([512, 1]))?;
        let bcast = model.wire_node(
            "bcast",
            MultiBroadcastTo {
                shape: ShapeFact::from_dims([1.to_dim(), 512.to_dim(), 16.to_dim()]),
            },
            &[src],
        )?[0];
        let unsq = model.wire_node("unsq", AxisOp::Add(3), &[bcast])?[0];
        model.select_output_outlets(&[unsq])?;

        let model = model.into_decluttered()?;
        assert_eq!(
            model.output_fact(0)?.shape.to_tvec(),
            tvec![1.to_dim(), 512.to_dim(), 16.to_dim(), 1.to_dim()]
        );
        Ok(())
    }

    /// An axis change touching an axis the input does not have must be
    /// absorbed in the target shape, not propagated to the input wire: a
    /// rank-1 [1] input cannot express it.
    #[test]
    fn broadcast_axis_change_absorbed_not_propagated_to_input() -> TractResult<()> {
        let mut model = TypedModel::default();
        let c = model.add_const("c", tensor1(&[0f32]))?;
        let s = model.symbols.sym("S");
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![s.to_dim(), 1usize.to_dim()])),
            &[c],
        )?[0];
        let node = model.node(y.node);
        let consequence = node
            .op
            .change_axes(&model, node, InOut::Out(0), &AxisOp::Add(2))?
            .context("expected the broadcast to absorb the change")?;
        assert_eq!(consequence.wire_changes, tvec![(InOut::Out(0), AxisOp::Add(2))]);
        Ok(())
    }

    /// Passthrough axis changes (input and output agree on the axis) still
    /// reach the input.
    #[test]
    fn passthrough_axis_change_propagates_to_input() -> TractResult<()> {
        let mut model = TypedModel::default();
        let s = model.symbols.sym("S");
        let src = model.add_source("src", f32::fact([s.to_dim(), 1usize.into()]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![s.to_dim(), 1usize.to_dim()])),
            &[src],
        )?[0];
        let node = model.node(y.node);
        let consequence = node
            .op
            .change_axes(&model, node, InOut::Out(0), &AxisOp::Add(1))?
            .context("expected propagation")?;
        assert_eq!(
            consequence.wire_changes,
            tvec![(InOut::Out(0), AxisOp::Add(1)), (InOut::In(0), AxisOp::Add(1))]
        );
        Ok(())
    }

    /// End-to-end through the ChangeAxes search: an axis change crossing a
    /// broadcast fed by a volume-1 constant must produce an applicable patch
    /// without any const rewiring.
    #[test]
    fn axis_change_through_broadcast_with_volume_one_const() -> TractResult<()> {
        let mut model = TypedModel::default();
        let s = model.symbols.sym("S");
        let definer = model.add_source("definer", f32::fact([s.to_dim()]))?;
        let c = model.add_const("c", tensor1(&[0f32]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![
                s.to_dim(),
                1usize.to_dim(),
                1usize.to_dim()
            ])),
            &[c],
        )?;
        let out = model.wire_node("out", AxisOp::Add(0), &[y[0]])?;
        model.select_output_outlets(&[out[0], definer])?;
        let change = AxisChange { outlet: out[0], op: AxisOp::Add(3) };
        let mut explored = Default::default();
        let (patch, _) =
            crate::optim::change_axes::change_axes(&model, &change, &[], &[], &mut explored)?
                .context("axis change through a broadcast should apply")?;
        patch.apply(&mut model)?;
        model.compact()?;
        let found = crate::internal::TypedSimplePlan::new(model)?
            .run(tvec!(tensor1(&[0f32; 2]).into_tvalue()))?;
        assert_eq!(found[0].shape(), &[1, 2, 1, 1, 1]);
        Ok(())
    }

    /// Removing an output axis on which the input is broadcast is absorbed:
    /// the right-aligned pairing is unaffected.
    #[test]
    fn rm_of_broadcast_axis_absorbed() -> TractResult<()> {
        let mut model = TypedModel::default();
        let src = model.add_source("src", f32::fact([5usize, 7usize]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![
                1usize.to_dim(),
                5usize.to_dim(),
                7usize.to_dim()
            ])),
            &[src],
        )?[0];
        let node = model.node(y.node);
        let consequence = node
            .op
            .change_axes(&model, node, InOut::Out(0), &AxisOp::Rm(0))?
            .context("expected the broadcast to absorb the removal")?;
        assert_eq!(consequence.wire_changes, tvec![(InOut::Out(0), AxisOp::Rm(0))]);
        let substitute = consequence
            .substitute_op
            .as_deref()
            .and_then(|op| op.as_op().downcast_ref::<MultiBroadcastTo>())
            .context("expected a MultiBroadcastTo substitute")?;
        assert_eq!(substitute.shape.to_tvec(), tvec![5.to_dim(), 7.to_dim()]);
        Ok(())
    }

    /// A Move touching an axis on which the input is broadcast must block:
    /// transpose-of-broadcast is not broadcast-into-transposed-target.
    #[test]
    fn move_over_broadcast_axis_blocked() -> TractResult<()> {
        let mut model = TypedModel::default();
        let src = model.add_source("src", f32::fact([5usize, 1usize]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![5usize.to_dim(), 7usize.to_dim()])),
            &[src],
        )?[0];
        let node = model.node(y.node);
        let blocked = node
            .op
            .change_axes(&model, node, InOut::Out(0), &AxisOp::Move(0, 1))
            .map(|r| r.is_none())
            .unwrap_or(false);
        assert!(blocked, "expected the broadcast to block the move");
        Ok(())
    }

    /// A passthrough Move still propagates to the input.
    #[test]
    fn passthrough_move_propagates_to_input() -> TractResult<()> {
        let mut model = TypedModel::default();
        let src = model.add_source("src", f32::fact([5usize, 7usize]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![5usize.to_dim(), 7usize.to_dim()])),
            &[src],
        )?[0];
        let node = model.node(y.node);
        let consequence = node
            .op
            .change_axes(&model, node, InOut::Out(0), &AxisOp::Move(0, 1))?
            .context("expected propagation")?;
        assert_eq!(
            consequence.wire_changes,
            tvec![(InOut::Out(0), AxisOp::Move(0, 1)), (InOut::In(0), AxisOp::Move(0, 1))]
        );
        Ok(())
    }

    /// The ChangeAxes pass proactively proposes removing any dim-1 output
    /// axis; the absorb must keep the input broadcastable into the new target
    /// ([5,1] into [9,5,1] must not become [9,5]).
    #[test]
    fn absorbed_change_keeps_input_broadcastable() -> TractResult<()> {
        let mut model = TypedModel::default();
        let src = model.add_source("src", f32::fact([5usize, 1usize]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![
                9usize.to_dim(),
                5usize.to_dim(),
                1usize.to_dim()
            ])),
            &[src],
        )?[0];
        let squeezed = model.wire_node("squeezed", AxisOp::Rm(2), &[y])?[0];
        model.select_output_outlets(&[squeezed])?;
        let optimized = model.into_optimized()?;
        let found = crate::internal::TypedSimplePlan::new(optimized)?
            .run(tvec!(tensor2(&[[0f32; 1]; 5]).into_tvalue()))?;
        assert_eq!(found[0].shape(), &[9, 5]);
        Ok(())
    }

    /// An absorbed change falling inside the input-paired window would
    /// re-pair the input's axes onto neighbouring target axes: with adjacent
    /// equal dims the result still evaluates, with permuted values. The
    /// optimized model must keep the unoptimized values.
    #[test]
    fn absorbed_change_does_not_repair_input_axes() -> TractResult<()> {
        let mut model = TypedModel::default();
        let src = model.add_source("src", f32::fact([5usize, 1usize]))?;
        let y = model.wire_node(
            "y",
            MultiBroadcastTo::new(ShapeFact::from_dims(tvec![
                5usize.to_dim(),
                5usize.to_dim(),
                1usize.to_dim()
            ])),
            &[src],
        )?[0];
        let squeezed = model.wire_node("squeezed", AxisOp::Rm(2), &[y])?[0];
        model.select_output_outlets(&[squeezed])?;
        let input: Vec<f32> = (0..5).map(|i| i as f32).collect();
        let input = Tensor::from_shape(&[5, 1], &input)?.into_tvalue();
        let raw =
            crate::internal::TypedSimplePlan::new(model.clone())?.run(tvec!(input.clone()))?;
        let optimized = model.into_optimized()?;
        let optimized_values =
            crate::internal::TypedSimplePlan::new(optimized)?.run(tvec!(input))?;
        assert_eq!(*raw[0], *optimized_values[0]);
        Ok(())
    }

    proptest! {
        /// The optimizer must preserve what the unoptimized model computes,
        /// whatever broadcast target and axis change it picks.
        #[test]
        fn prop_axis_changes_preserve_values(
            input in proptest::collection::vec(1usize..4, 1usize..4),
            leading in proptest::collection::vec(1usize..4, 0usize..3),
            bumps in proptest::collection::vec(0usize..4, 0usize..6),
            selector in 0usize..64,
            ix in 0usize..8,
            move_to in 0usize..8,
        ) {
            let mut bumps = bumps.into_iter();
            let mut target: TVec<usize> = leading.clone().into();
            for dim in &input {
                let dim = if *dim == 1 { bumps.next().unwrap_or(0).max(1) } else { *dim };
                target.push(dim);
            }
            let op = match selector % 3 {
                0 => AxisOp::Add(ix.min(target.len())),
                1 => {
                    if ix < target.len() && target[ix] == 1 {
                        AxisOp::Rm(ix)
                    } else {
                        AxisOp::Add(ix.min(target.len()))
                    }
                }
                _ => {
                    let a = ix.min(target.len().saturating_sub(1));
                    let b = (move_to.min(target.len().saturating_sub(1)), a);
                    let b = if b.0 == a { a.saturating_sub(1) } else { b.0 };
                    AxisOp::Move(a, b)
                }
            };
            let mut model = TypedModel::default();
            let src = model.add_source("src", i32::datum_type().fact(&*input)).unwrap();
            let bcast = model.wire_node(
                "bcast",
                MultiBroadcastTo::new(ShapeFact::from_dims(
                    target.iter().map(|d| d.to_dim()).collect::<TVec<_>>(),
                )),
                &[src],
            ).unwrap()[0];
            let out = model.wire_node("out", op, &[bcast]).unwrap()[0];
            model.select_output_outlets(&[out]).unwrap();
            let volume: usize = input.iter().product();
            let data: Vec<i32> = (0..volume as i32).collect();
            let wire = Tensor::from_shape(&input, &data).unwrap().into_tvalue();
            let raw = crate::internal::TypedSimplePlan::new(model.clone())
                .unwrap()
                .run(tvec!(wire.clone()))
                .unwrap();
            let optimized = model.into_optimized().unwrap();
            let optimized_values = crate::internal::TypedSimplePlan::new(optimized)
                .unwrap()
                .run(tvec!(wire))
                .unwrap();
            prop_assert_eq!(&*raw[0], &*optimized_values[0]);
        }
    }
}