automatic_differentiation.py

"""
Demonstration of the Automatic Differentiation (Reverse mode).

Reference: https://en.wikipedia.org/wiki/Automatic_differentiation

Author: Poojan Smart
Email: smrtpoojan@gmail.com
"""

from __future__ importannotations

fromcollectionsimportdefaultdict
fromenumimportEnum
fromtypesimportTracebackType
fromtypingimportAny

importnumpyasnp
fromtyping_extensionsimportSelf# noqa: UP035


classOpType(Enum):
"""
 Class represents list of supported operations on Variable for gradient calculation.
 """

ADD=0
SUB=1
MUL=2
DIV=3
MATMUL=4
POWER=5
NOOP=6


classVariable:
"""
 Class represents n-dimensional object which is used to wrap numpy array on which
 operations will be performed and the gradient will be calculated.

 Examples:
 >>> Variable(5.0)
 Variable(5.0)
 >>> Variable([5.0, 2.9])
 Variable([5. 2.9])
 >>> Variable([5.0, 2.9]) + Variable([1.0, 5.5])
 Variable([6. 8.4])
 >>> Variable([[8.0, 10.0]])
 Variable([[ 8. 10.]])
 """

def__init__(self, value: Any) ->None:
self.value=np.array(value)

# pointers to the operations to which the Variable is input
self.param_to: list[Operation] = []
# pointer to the operation of which the Variable is output of
self.result_of: Operation=Operation(OpType.NOOP)

def__repr__(self) ->str:
returnf"Variable({self.value})"

defto_ndarray(self) ->np.ndarray:
returnself.value

def__add__(self, other: Variable) ->Variable:
result=Variable(self.value+other.value)

withGradientTracker() astracker:
# if tracker is enabled, computation graph will be updated
iftracker.enabled:
tracker.append(OpType.ADD, params=[self, other], output=result)
returnresult

def__sub__(self, other: Variable) ->Variable:
result=Variable(self.value-other.value)

withGradientTracker() astracker:
# if tracker is enabled, computation graph will be updated
iftracker.enabled:
tracker.append(OpType.SUB, params=[self, other], output=result)
returnresult

def__mul__(self, other: Variable) ->Variable:
result=Variable(self.value*other.value)

withGradientTracker() astracker:
# if tracker is enabled, computation graph will be updated
iftracker.enabled:
tracker.append(OpType.MUL, params=[self, other], output=result)
returnresult

def__truediv__(self, other: Variable) ->Variable:
result=Variable(self.value/other.value)

withGradientTracker() astracker:
# if tracker is enabled, computation graph will be updated
iftracker.enabled:
tracker.append(OpType.DIV, params=[self, other], output=result)
returnresult

def__matmul__(self, other: Variable) ->Variable:
result=Variable(self.value @ other.value)

withGradientTracker() astracker:
# if tracker is enabled, computation graph will be updated
iftracker.enabled:
tracker.append(OpType.MATMUL, params=[self, other], output=result)
returnresult

def__pow__(self, power: int) ->Variable:
result=Variable(self.value**power)

withGradientTracker() astracker:
# if tracker is enabled, computation graph will be updated
iftracker.enabled:
tracker.append(
OpType.POWER,
params=[self],
output=result,
other_params={"power": power},
 )
returnresult

defadd_param_to(self, param_to: Operation) ->None:
self.param_to.append(param_to)

defadd_result_of(self, result_of: Operation) ->None:
self.result_of=result_of


classOperation:
"""
 Class represents operation between single or two Variable objects.
 Operation objects contains type of operation, pointers to input Variable
 objects and pointer to resulting Variable from the operation.
 """

def__init__(
self,
op_type: OpType,
other_params: dict|None=None,
 ) ->None:
self.op_type=op_type
self.other_params= {} ifother_paramsisNoneelseother_params

defadd_params(self, params: list[Variable]) ->None:
self.params=params

defadd_output(self, output: Variable) ->None:
self.output=output

def__eq__(self, value) ->bool:
returnself.op_type==valueifisinstance(value, OpType) elseFalse


classGradientTracker:
"""
 Class contains methods to compute partial derivatives of Variable
 based on the computation graph.

 Examples:

 >>> with GradientTracker() as tracker:
 ... a = Variable([2.0, 5.0])
 ... b = Variable([1.0, 2.0])
 ... m = Variable([1.0, 2.0])
 ... c = a + b
 ... d = a * b
 ... e = c / d
 >>> tracker.gradient(e, a)
 array([-0.25, -0.04])
 >>> tracker.gradient(e, b)
 array([-1. , -0.25])
 >>> tracker.gradient(e, m) is None
 True

 >>> with GradientTracker() as tracker:
 ... a = Variable([[2.0, 5.0]])
 ... b = Variable([[1.0], [2.0]])
 ... c = a @ b
 >>> tracker.gradient(c, a)
 array([[1., 2.]])
 >>> tracker.gradient(c, b)
 array([[2.],
 [5.]])

 >>> with GradientTracker() as tracker:
 ... a = Variable([[2.0, 5.0]])
 ... b = a ** 3
 >>> tracker.gradient(b, a)
 array([[12., 75.]])
 """

instance=None

def__new__(cls) ->Self:
"""
 Executes at the creation of class object and returns if
 object is already created. This class follows singleton
 design pattern.
 """
ifcls.instanceisNone:
cls.instance=super().__new__(cls)
returncls.instance

def__init__(self) ->None:
self.enabled=False

def__enter__(self) ->Self:
self.enabled=True
returnself

def__exit__(
self,
exc_type: type[BaseException] |None,
exc: BaseException|None,
traceback: TracebackType|None,
 ) ->None:
self.enabled=False

defappend(
self,
op_type: OpType,
params: list[Variable],
output: Variable,
other_params: dict|None=None,
 ) ->None:
"""
 Adds Operation object to the related Variable objects for
 creating computational graph for calculating gradients.

 Args:
 op_type: Operation type
 params: Input parameters to the operation
 output: Output variable of the operation
 """
operation=Operation(op_type, other_params=other_params)
param_nodes= []
forparaminparams:
param.add_param_to(operation)
param_nodes.append(param)
output.add_result_of(operation)

operation.add_params(param_nodes)
operation.add_output(output)

defgradient(self, target: Variable, source: Variable) ->np.ndarray|None:
"""
 Reverse accumulation of partial derivatives to calculate gradients
 of target variable with respect to source variable.

 Args:
 target: target variable for which gradients are calculated.
 source: source variable with respect to which the gradients are
 calculated.

 Returns:
 Gradient of the source variable with respect to the target variable
 """

# partial derivatives with respect to target
partial_deriv=defaultdict(lambda: 0)
partial_deriv[target] =np.ones_like(target.to_ndarray())

# iterating through each operations in the computation graph
operation_queue= [target.result_of]
whilelen(operation_queue) >0:
operation=operation_queue.pop()
forparaminoperation.params:
# as per the chain rule, multiplying partial derivatives
# of variables with respect to the target
dparam_doutput=self.derivative(param, operation)
dparam_dtarget=dparam_doutput*partial_deriv[operation.output]
partial_deriv[param] +=dparam_dtarget

ifparam.result_ofandparam.result_of!=OpType.NOOP:
operation_queue.append(param.result_of)

returnpartial_deriv.get(source)

defderivative(self, param: Variable, operation: Operation) ->np.ndarray:
"""
 Compute the derivative of given operation/function

 Args:
 param: variable to be differentiated
 operation: function performed on the input variable

 Returns:
 Derivative of input variable with respect to the output of
 the operation
 """
params=operation.params

ifoperation==OpType.ADD:
returnnp.ones_like(params[0].to_ndarray(), dtype=np.float64)
ifoperation==OpType.SUB:
ifparams[0] ==param:
returnnp.ones_like(params[0].to_ndarray(), dtype=np.float64)
return-np.ones_like(params[1].to_ndarray(), dtype=np.float64)
ifoperation==OpType.MUL:
return (
params[1].to_ndarray().T
ifparams[0] ==param
elseparams[0].to_ndarray().T
 )
ifoperation==OpType.DIV:
ifparams[0] ==param:
return1/params[1].to_ndarray()
return-params[0].to_ndarray() / (params[1].to_ndarray() **2)
ifoperation==OpType.MATMUL:
return (
params[1].to_ndarray().T
ifparams[0] ==param
elseparams[0].to_ndarray().T
 )
ifoperation==OpType.POWER:
power=operation.other_params["power"]
returnpower* (params[0].to_ndarray() ** (power-1))

err_msg=f"invalid operation type: {operation.op_type}"
raiseValueError(err_msg)


if__name__=="__main__":
importdoctest

doctest.testmod()