The Good and The Evil Way to Perform Monkey Patching

[cgarciae]

Here we will show two ways to monkey patch an object to get new behavior. As inspiration lets first create a simple model.

from flax import nnx
import jax.numpy as jnp

class Foo(nnx.Module):
  def __init__(self, rngs):
    self.linear1 = nnx.Linear(2, 2, rngs=rngs)
    self.linear2 = nnx.Linear(2, 2, rngs=rngs)
  
  def __call__(self, x):
    return self.linear2(nnx.relu(self.linear1(x)))
    
model = Foo(nnx.Rngs(0))
x = jnp.ones((1, 2))

Now lets say we want to update linear1 to offset its output by 100. How can we do this?

The Good Way

The first way is to create a new type that inherits from Linear and override its __call__ method to perform the new behavior. We also override the constructor to be able to easily create it from an existing instance by copying all its attributes. Then we simply swap model.linear1 with its new value.

class CrazyLinear(nnx.Linear):
  def __init__(self, existing):
    vars(self).update(vars(existing)) # copy over state
  
  def __call__(self, x): # override method
    return super().__call__(x) + 100
    
model.linear1 = CrazyLinear(model.linear1)  # swap object

y = model(x)

The Evil Way

Another way of doing this is to create a patch_method function that can take any object and overrides any given method with a new provided method. It does so by 1) programmatically creating a new type that inherits from the existing type and overrides the target method with the new method, and 2) updates the object type in place‼️

Then you can simply define the new method and use patch_method to update the __call__ method for model.linear1.

def patch_method(obj, name, new_method):
  # create new type and change object's type 😈
  obj.__class__ = type(
    obj.__class__.__name__,
    (obj.__class__,),
    {name: new_method},
  )
  
def crazy_linear_call(self, x):
  return nnx.Linear.__call__(self, x) + 100
  
patch_method(model.linear1, '__call__', crazy_linear_call)

y = model(x)

Warning: use with care.