pycodon_autograd_01

import math
import time
from typing import List, Optional

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

class NaiveVar:
    val: float
    grad: float

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

class MyOperation:
    op_type: int
    n_inputs: int
    inputs: List[NaiveVar]
    outputs: List[NaiveVar]

    def __init__(self, op_type: int, n_inputs: int, inputs: List[NaiveVar], outputs: List[NaiveVar]) -> None:
        self.op_type = op_type
        self.n_inputs = n_inputs
        self.inputs = inputs
        self.outputs = outputs

class NaiveTape:
    n_ops: int
    max_ops: int
    ops: List[Optional[MyOperation]]

    def __init__(self, max_ops: int) -> None:
        self.n_ops = 0
        self.max_ops = max_ops
        self.ops = [None] * max_ops

    def add_op(self, op: MyOperation) -> None:
        if self.n_ops >= self.max_ops:
            self.max_ops *= 2
            self.ops.extend([None] * (self.max_ops - len(self.ops)))
        self.ops[self.n_ops] = op
        self.n_ops += 1

def naive_var(val: float) -> NaiveVar:
    return NaiveVar(val)

def naive_tape_create(max_ops: int) -> NaiveTape:
    return NaiveTape(max_ops)

def naive_tape_sum(tape: NaiveTape, vars: List[NaiveVar]) -> NaiveVar:
    sum_val = sum(var.val for var in vars)
    res = naive_var(sum_val)
    op = MyOperation(OpType.OP_SUM, len(vars), vars, [res])
    tape.add_op(op)
    return res

def naive_tape_prod(tape: NaiveTape, var1: NaiveVar, var2: NaiveVar) -> NaiveVar:
    prod_val = var1.val * var2.val
    res = naive_var(prod_val)
    op = MyOperation(OpType.OP_PROD, 2, [var1, var2], [res])
    tape.add_op(op)
    return res

def naive_tape_softmax(tape: NaiveTape, vars: List[NaiveVar]) -> List[NaiveVar]:
    vals = [var.val for var in vars]
    maxval = max(vals)
    vals = [val - maxval for val in vals]
    denom = sum(math.exp(val) for val in vals)
    res = [naive_var(math.exp(val) / denom) for val in vals]
    op = MyOperation(OpType.OP_SOFTMAX, len(vars), vars, res)
    tape.add_op(op)
    return res

def naive_tape_softplus(tape: NaiveTape, var: NaiveVar) -> NaiveVar:
    softplus_val = math.log1p(math.exp(var.val))
    res = naive_var(softplus_val)
    op = MyOperation(OpType.OP_SOFTPLUS, 1, [var], [res])
    tape.add_op(op)
    return res

def naive_tape_backward(tape: NaiveTape, var: NaiveVar) -> None:
    var.grad = 1.0  # Set gradient of the output variable
    for i in range(tape.n_ops - 1, -1, -1):
        op = tape.ops[i]
        if op is None:
            continue
        if op.op_type == OpType.OP_SUM:
            for j in range(op.n_inputs):
                op.inputs[j].grad += op.outputs[0].grad
        elif op.op_type == OpType.OP_PROD:
            in1 = op.inputs[0]
            in2 = op.inputs[1]
            in1.grad += in2.val * op.outputs[0].grad
            in2.grad += in1.val * op.outputs[0].grad
        elif op.op_type == OpType.OP_SOFTPLUS:
            op.inputs[0].grad += op.outputs[0].grad / (1.0 + math.exp(-op.inputs[0].val))
        elif op.op_type == OpType.OP_SOFTMAX:
            # Softmax backward not implemented for now
            pass
        else:
            # Handle unknown operations
            pass

def main() -> None:
    start_time = time.time()
    iterations = 1000000

    for i in range(iterations):
        # Create variables
        x = naive_var(1.0)
        y = naive_var(2.0)

        # Create tape
        tape = naive_tape_create(10)

        # Perform operations
        sum_xy = naive_tape_sum(tape, [x, y])
        prod_sum_xy = naive_tape_prod(tape, sum_xy, sum_xy)
        softplus_prod = naive_tape_softplus(tape, prod_sum_xy)

        # Backward pass
        naive_tape_backward(tape, softplus_prod)

        if i == iterations - 1:
            print(f"sum_xy.val: {sum_xy.val}")
            print(f"prod_sum_xy.val: {prod_sum_xy.val}")
            print(f"softplus_prod.val: {softplus_prod.val}")
            print(f"sum_xy.grad: {sum_xy.grad}")
            print(f"x.grad: {x.grad}")
            print(f"y.grad: {y.grad}")

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


main()