How does grad treat self?

[dieterichlawson]

I was playing around with trying to differentiate functions that accept a self parameter because they're defined on objects. I was expecting grad to fail in this case, but it seems to ignore self entirely. Here is some example code that is trying to compute the gradient of sum(W @ x) with respect to W:

class Linear:

  W: jnp.ndarray

  def __init__(
          self, in_dim: int, out_dim: int, key: PRNGKey,
          W_init=jax.nn.initializers.glorot_normal()):
    self.W = W_init(key, (out_dim, in_dim))

  def __call__(self, x):
    return jnp.sum(self.W @ x)


lin = Linear(3, 2, jax.random.PRNGKey(0))

x = jax.random.uniform(jax.random.PRNGKey(0), (3,))
print(jax.grad(lin)(x))
print(jnp.sum(lin.W, axis=0))

grad does not see self as an argument and instead treats x as the first argument, resulting in it returning the gradient of sum(W @ x) with respect to x. What exactly is going on here? I looked in the code for grad and value_and_grad and was unable to figure out where this behavior was implemented.

[YouJiacheng]

Jax use tracing.
Callable object just works like a closure, namely lin(x) == (lambda x: Linear.__call__(lin, x))(x).