Skip to content

Custom (mutable) traversal is a bit hard. #159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
the-ssd opened this issue Apr 16, 2025 · 8 comments · Fixed by #160
Closed

Custom (mutable) traversal is a bit hard. #159

the-ssd opened this issue Apr 16, 2025 · 8 comments · Fixed by #160

Comments

@the-ssd
Copy link

the-ssd commented Apr 16, 2025

Currently, the best I was able to do is passing the tree as mutable, and NodeIdx

getting children is also a bit weird

let node = self.tree.get_node(&node_idx).unwrap();
let children_ids = node.children().map(|child| child.idx()).collect::<Vec<_>>();

let mut children = vec![];
for node in children_ids {
    let value = self.execute_rec(inputs, node);
    children.push(value);
}

let mut node = self.tree.get_node_mut(&node_idx).unwrap();
// Use children and node here

That being said, this is still a very useful create.

@orxfun
Copy link
Owner

orxfun commented Apr 17, 2025

Thanks a lot for the feedback @the-ssd , it is very helpful!

Indeed, mutable traversal is the most complicated one. But hope we can fix the issue and find a convenient way.

I created this branch (https://github.com/orxfun/orx-tree/tree/mutable-traversal) and added an example (https://github.com/orxfun/orx-tree/blob/mutable-traversal/examples/mutable_traversal.rs) there.

Currently, there are two versions:

  • implementation1: I tried to keep this as close as your implementation above.
  • implementation2: Not so much different, only uses node.children()

Then, I noticed that the important bit is // Use children and node here part of the code. It would be great

  • if you could add the continuation of the code block above (doesn't necessarily need to compile),
  • or feel free to directly commit to the example,
  • or feel free to link me to the source code, might help me understand better.

I believe the solution might depend on this part.

@the-ssd
Copy link
Author

the-ssd commented Apr 17, 2025

My uses might be a little weird, but here it is:

I have an Instruction enum:

enum Intruction {
    Input,
    Add,
    AddI(f32)
}

And the tree is a DynTree<(Instruction, f32)>,

Basically, I am interpreting some code, and storing the result of each step in the tree to use later.

.children() Api almost worked for be, but I wasn't able to start the recusive loop because I needed a NodeMutDown instead of NodeMutUpAndDown that I get from tree.root_mut()

Here is the main part:

fn execute_rec(
    &mut self,
    inputs: &[f32; INPUTS_COUNT],
    node_idx: NodeIdx<Dyn<(Instruction, f32)>>,
) -> f32 {
    let node = self.tree.get_node(&node_idx).unwrap();

    let children_ids = node.children().map(|child| child.idx()).collect::<Vec<_>>();
    let mut children = vec![];

    for node in children_ids {
        let value = self.execute_rec(inputs, node);
        children.push(value);
    }

    let mut node = self.tree.get_node_mut(&node_idx).unwrap();

    let (inst, value) = node.data_mut();

    *value = match inst {
        Instruction::Input(i) => inputs[*i],

        Instruction::Add => children.into_iter().sum(),
        // ...

        Instruction::AddI { val } => children.into_iter().sum::<f32>() + *val,
        // ...
    };

    *value
}

Maybe adding a way to convert UpAndDown Nodes into only Down might solve it.

An example can be writing something like an inverse depth to each node:
recursively descend the tree for leaves write 0, for other nodes write highest (or smallest) of its children + 1

@orxfun
Copy link
Owner

orxfun commented Apr 17, 2025

Thank you! This is a very interesting use case.

I have updated the example wrt your feedback.

It would be great if you could test it out:

cargo run --example mutable_traversal

There are again two implementations:

  • impl_over_children_idx as close as possible to what you have provided, only made some changes to be able to Display
  • impl_over_children_mut is the alternative way where we do not need to collect indices and results. This calls the following alternative recursive execute function
fn execute<'a>(
    inputs: &[f32; INPUTS_COUNT],
    mut node: NodeMut<'a, Dyn<InstructionNode>>,
) -> (NodeMut<'a, Dyn<InstructionNode>>, f32) {
    let num_children = node.num_children();

    let new_value = match node.data().instruction {
        Instruction::Input(i) => inputs[i],
        Instruction::Add => {
            let mut value = 0.0;
            for i in 0..num_children {
                let child = node.into_child_mut(i).unwrap();
                let (child, child_value) = execute(inputs, child);
                value += child_value;
                node = child.into_parent_mut().unwrap();
            }
            value
        }
        Instruction::AddI { val } => {
            let mut value = val;
            for i in 0..num_children {
                let child = node.into_child_mut(i).unwrap();
                let (child, child_value) = execute(inputs, child);
                value += child_value;
                node = child.into_parent_mut().unwrap();
            }
            value
        }
    };

    (*node.data_mut()).value = new_value;

    (node, new_value)
}

This is a bit tricky but a safe way to traversal.

  • We may freely move up & down.
  • We may mutate the shape of the tree while doings so.
  • We may mutate the node's data.

These are all safe because we always only have one mutable node. We keep converting our mutable node to child and to parent.

This second implementation has the advantage of being truly recursive and not requiring intermediate collections. But there is currently a convenience problem. As apparent in two branches in the match within the execute method, we need to hand write the traversal.

  • (i) This is one thing that can improve. We can abstract the depth-first traversal with mutable references, accepting a closure to update the node data.

However, as you also noted, we might also have an alternative way. It might actually be safe in sequential code to walk down the tree, say in bfs order, while we only allow to mutate data but not the structure of the tree. Currently, we can create a traversal or walk from any node where elements can be the data or the node itself which would allow for our recursion.

let node = self.tree.node(&node_idx);
let mut dfs = Traversal.dfs().over_data();
for data in node.walk_with(&mut dfs) {
    // data: &InstructionNode>
}

let node = self.tree.node(&node_idx);
let mut dfs = Traversal.dfs().over_nodes();
for node in node.walk_with(&mut dfs) {
    // node: Node<'_, Dyn<InstructionNode>>
}

However, we can create a mutable traversal only on the data.

let mut node_mut = self.tree.node_mut(&node_idx);
let mut dfs = Traversal.dfs().over_data();
for data_mut in node_mut.walk_mut_with(&mut dfs) {
    // data_mut: &mut InstructionNode>
}

let mut node_mut = self.tree.node_mut(&node_idx);
let mut dfs = Traversal.dfs().over_nodes();
for node_mut in node_mut.walk_mut_with(&mut dfs) {
    // node_mut: NodeMut<'a, Dyn<InstructionNode>>
    // does not compile right now!
    // with this, we could've made the recursive call
}

(ii) I believe this is also possible but we need to implement another NodeMutOrientation which:

  • allows to update data
  • but not the shape, push_child should not be available for instance.

I will work on these two points (i) and (ii). But will be on vacation for a couple of weeks.

Meanwhile, please let me know if you have other suggestions.

@orxfun
Copy link
Owner

orxfun commented Apr 17, 2025

Actually, I think there was a bug in the execute function above; regardless of the instruction variant, we must reach the leaves to compute all nodes. And now it looks nicer :)

fn execute<'a>(
    inputs: &[f32; INPUTS_COUNT],
    mut node: NodeMut<'a, Dyn<InstructionNode>>,
) -> (NodeMut<'a, Dyn<InstructionNode>>, f32) {
    let num_children = node.num_children();

    let mut children_sum = 0.0;
    for i in 0..num_children {
        let child = node.into_child_mut(i).unwrap();
        let (child, child_value) = execute(inputs, child);
        children_sum += child_value;
        node = child.into_parent_mut().unwrap();
    }

    let new_value = match node.data().instruction {
        Instruction::Input(i) => inputs[i],
        Instruction::Add => children_sum,
        Instruction::AddI { val } => val + children_sum,
    };

    (*node.data_mut()).value = new_value;

    (node, new_value)
}

@orxfun
Copy link
Owner

orxfun commented Apr 26, 2025

@the-ssd , thank you for bringing this interesting use case. The PR #160 aims to resolve this issue in two different ways:

  • it introduces the recursive_set method which provides an expressive way to handle this task while abstracting away the traversal and recursion,
  • it presents an example based on your use case where different approaches are demonstrated and documented with their pros and cons.

Please let me know of your feedback.

@the-ssd
Copy link
Author

the-ssd commented Apr 26, 2025

This looks perfect! Thank you for this awesome crate.

@the-ssd the-ssd closed this as completed Apr 26, 2025
@the-ssd
Copy link
Author

the-ssd commented Apr 26, 2025

Also, while this is a bit out of scope, I think it might make sense to explain why I wanted a tree instead of a Vec.

So there are programs called superoptimizers. Basically like a normal optimizer, but instead of transforming a program into a faster one, it synthesizes a program that is (probably) the fastest possible.

There are a few way to do that. Some use brute force and try all programs. Some use a SAT solver, and some like stoke use stochastic optimization (basically mutate a program and sometimes accept a worse program). Their paper on stoke is really interested in my opinion (they were able to find a program that is better than expert written Assembly!), and I even (kind of) was able to replicate it. I used it only on some samples of smoothstep (without clamping), and it was able to synthesizes a smoothstep function, and even use 1 less instruction than what I expected!

Then I was trying to do something similar to backpropagation, and for this a tree is much easier than a Vec, because 2 instruction might want different changes from previous state, which isn't the case with a tree.

In the end, I lost motivation, and didn't get anywhere. I thought it might make sense since you seem interested in similar things.

@orxfun
Copy link
Owner

orxfun commented Apr 27, 2025

Thanks the-ssd. Indeed, these are very interesting topics for me. I am an operations research scientist by day mostly trying to optimize real life problems. Actually my goal is to converge to optimization libraries in my domain, although I am side-stepping a bit with preliminaries :) which is also fun.

Superoptimizers are completely new to me but found the idea amazing, and stochastic optimization approach seems very promising at first glance. I will certainly read more on this. Thanks again for sharing. In the meantime, feel free to open an issue whenever needed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants