Tree or Graph based data structure

[sa-]

Hey fellow Mojovians Mojicians

I'm trying to implement backpopogation for a neural net and need to represent a computational graph somehow and the easiest way to do it is through some sort of tree.

struct Node:
    var left: Self
    var right: Self

Now I can't add a @value decorator because Node is a non-copyable and non-movable type. And if I create a construtor I wouldn't know what to set left or right as because we don't have null values yet.

And the kernel crashes when I run

@value
struct Node:
    var left: Pointer[Self]
    var right: Pointer[Self]

I can't go with an adjacency matrix either because we don't have Dict types yet, it feels like I've hit a roadblock. Does anyone have ideas/workarounds for storing graph/tree like data structures?

The above snippet is a reduced version, this is what my Tensor struct looks like

from Vector import DynamicVector
struct Tensor:
    var shape: DynamicVector[Int] # TODO use type param for type safe ops
    var data: DynamicVector[F64]
    var operation: Int
    var gradient: F64
    var child1: Self
    var child2: Self
    
    alias OP_VOID = 0
    alias OP_MULTIPLY = 1
    alias OP_ADD = 2
    alias OP_SUBTRACT = 3
    alias OP_POW = 4
    alias OP_EXP = 5

   ...

When Tensors are initialised I only want shape and data to be defined but not the children,
and when I call let t3 = t1.add(t2) I would like the add function to set t3.child1 as t1, and t3.child2 as t2

[lattner]

The solution with Pointer should work, but you'll need to be /extremely/ careful with them - in the absence of lifetimes, it is super easy to introduce dangling pointers etc. You'll also have to manually implement the copy/move ctors.

btw, I'm fond of mojician 🪄

[sa-]

First of all, Mojician sounds great!

I'm going to think out loud now-
Another could be to create just the computation graph (or lazy tensors) and to basically allocate a massive fixed size Vector upfront at calculation time, since we would know the count and sizes of the Tensors.

This would be safer (no need for pointers) and maybe even more performant since I'm pretty certain we won't need to do heap allocations anymore. But I also only have until Tuesday until my job starts again, this sounds like more than a day's effort, and I won't be able to dedicate as much time here after that.

[sa-]

Follow up, the pointer version seems to throw an error

@value
struct Node:
    var left: Pointer[Self]
    var right: Pointer[Self]

let n = Node(Pointer[Node].get_null(), Pointer[Node].get_null())

Error:

error: Expression [2]:16:5: 'pop.store' op failed to verify that pointer element type
    let n = Node(Pointer[Node].get_null(), Pointer[Node].get_null())
    ^

Edit: Made #291 for this

[lattner]

@Mogball FYI, likely the same issue as #291 (got split off)

[cjohnson318]

Mojician FTW

[automata]

I'm trying to port micrograd to Mojo while learning the language, and facing exactly the same problem. The following snippet crashes the kernel:

from String import String
from BuiltinList import FloatLiteral

struct Value:
    var data: Int
    var grad: FloatLiteral
    var _prev_a: Self
    var _prev_b: Self
    var _op: String

    fn __init__(inout self, data: Int, _child_a: Self, _child_b: Self, _op: String):
        self.data = data
        self.grad = 0.0
        self._prev_a = _child_a
        self._prev_b = _child_b
        self._op = _op

    fn __copyinit__(inout self, other: Self):
        self.data = other.data
        self.grad = other.grad
        self._prev_a = other._prev_a
        self._prev_b = other._prev_b
        self._op = other._op

    fn __add__(self, owned other: Self) -> Self:
        let out: Value = Value(self.data + other.data, self, other, '+')
        return out

    fn _backward_add(inout self, inout other: Self):
        self._prev_a.grad += self.grad
        self._prev_b.grad += self.grad

I was trying to isolate the problem while following your suggestion to use Pointers. However, while I'm able to do:

var a: Pointer[Int]
a = Pointer[Int].alloc(1)
a.store(0, 42)
print(a.load(0))

I'm not able to do the same with:

struct Node:
    var left: Pointer[Self]
    var right: Pointer[Self]
    
    fn __init__(inout self):
        self.left = Pointer[Self].alloc(1)
        self.right = Pointer[Self].alloc(1)

Getting the following error as result:

error: Expression [20]:11:40: symbol use result #0 has type '!pop.pointer<@"$Expression [20]::Node">' but @$Pointer::Pointer::alloc($Int::Int) expected type '!pop.pointer<struct<pointer<@"$Expression [20]::Node">, pointer<@"$Expression [20]::Node">>>'
        self.left = Pointer[Node].alloc(1)
                                       ^

/.modular/Kernels/mojo/Stdlib/Pointer.mojo:186:5: @$Pointer::Pointer::alloc($Int::Int) declared here
    fn alloc(count: Int) -> Pointer[type]:
    ^

I'm pretty new to Mojo and MLIR, so there's a big chance that I'm using the pointers in the wrong way 😅 Any suggestion is really welcomed!

[sa-]

Hey @automata , nice to see someone else working in this domain.

You can see my comment above for the code snippet that I expect should work and the corresponding error message

[automata]

Hey @sa-, really great to find you and all you're doing around the Tensor type, I really like it!

Yeah, thank you for pointing me out for that, curious to see if that's actually a bug.

[RWayne93]

@automata were you able to get a micrograd port working in mojo? i too had tried it and was crashing my kernel as well in the playground environment.