codon_autograd_v02a

from typing import List, Tuple
import math
import random
from time import time

MAX_VARS = 1000
MAX_OPS = 1000
MAX_INPUTS = 2


class OpType:
    OP_SUM = 0
    OP_PROD = 1
    OP_SOFTPLUS = 2


class NaiveVar:
    val: float
    grad: float

    def __init__(self, val: float) -> None:
        self.val = val
        self.grad = 0.0


class Operation:
    op_type: int
    input_ids: Tuple[int, int]
    output_id: int

    def __init__(self, op_type: int, input_ids: Tuple[int, int], output_id: int) -> None:
        self.op_type = op_type
        self.input_ids = input_ids
        self.output_id = output_id


class NaiveTape:
    vars: List[NaiveVar]
    ops: List[Operation]

    def __init__(self):
        self.vars = []
        self.ops = []


    def create_var(self, val: float) -> int:
        var = NaiveVar(val)
        id = len(self.vars)
        self.vars.append(var)
        return id

    def add_operation(self, op: Operation) -> None:
        self.ops.append(op)

    def backward(self, var_id: int) -> None:
        self.vars[var_id].grad = 1.0
        i = len(self.ops)

        while i > 0:
            i -= 1
            op = self.ops[i]
            output_grad = self.vars[op.output_id].grad

            if op.op_type == OpType.OP_SUM:
                self.vars[op.input_ids[0]].grad += output_grad
                self.vars[op.input_ids[1]].grad += output_grad

            elif op.op_type == OpType.OP_PROD:
                input1_val = self.vars[op.input_ids[0]].val
                input2_val = self.vars[op.input_ids[1]].val
                self.vars[op.input_ids[0]].grad += input2_val * output_grad
                self.vars[op.input_ids[1]].grad += input1_val * output_grad

            elif op.op_type == OpType.OP_SOFTPLUS:
                input_val = self.vars[op.input_ids[0]].val
                exp_val = math.exp(-input_val)
                self.vars[op.input_ids[0]].grad += output_grad / (1.0 + exp_val)



def main():
    iterations: int = 1000000
    # random.seed(time.time())
    start_time = time()

    for i in range(iterations):
        tape = NaiveTape()
        var1_id = tape.create_var(random.random())
        var2_id = tape.create_var(random.random())

        # Inlined sum
        sum_val = tape.vars[var1_id].val + tape.vars[var2_id].val
        sum_var_id = tape.create_var(sum_val)
        op_sum = Operation(OpType.OP_SUM, (var1_id, var2_id), sum_var_id)
        tape.add_operation(op_sum)

        # Inlined prod
        prod_val = tape.vars[sum_var_id].val * tape.vars[sum_var_id].val
        prod_var_id = tape.create_var(prod_val)
        op_prod = Operation(OpType.OP_PROD, (sum_var_id, sum_var_id), prod_var_id)
        tape.add_operation(op_prod)

        # Inlined softplus
        softplus_val = math.log1p(math.exp(tape.vars[prod_var_id].val))
        softplus_var_id = tape.create_var(softplus_val)
        op_softplus = Operation(OpType.OP_SOFTPLUS, (prod_var_id, 0), softplus_var_id)
        tape.add_operation(op_softplus)

        tape.backward(softplus_var_id)

        if i == iterations - 1:
            print(f"sum_var val: {tape.vars[sum_var_id].val}")
            print(f"prod_var val: {tape.vars[prod_var_id].val}")
            print(f"softplus_var val: {tape.vars[softplus_var_id].val}")
            print(f"sum_var grad: {tape.vars[sum_var_id].grad}")
            print(f"var1 grad: {tape.vars[var1_id].grad}")
            print(f"var2 grad: {tape.vars[var2_id].grad}")

    end_time = time()
    elapsed_time = (end_time - start_time) * 1000
    print(f"\nElapsed time: {elapsed_time:.3f} ms")


main()