Trying to encapsulate jax code into a class

[tw40210]

Description
I want to encapsulate my jax code into a class to call the member function update itself. However, the update function in the class only update once and don't change anymore. I've read some material and example of jax but they usually are written out of a class. Thus, I can't find one to fit my needs. Please help me to figure out what I did wrongly. Thank you!!

What jax/jaxlib version are you using?
jax 0.2.19
jaxlib 0.1.70

Which accelerator(s) are you using?
CPU

Additional system info
Windows WSL

NVIDIA GPU info
No response

My Code

import jax.numpy as jnp
import jax.tree_util as jtree
import jax.random as jrand
import matplotlib.pyplot as plt
from functools import partial

def JaxFC(n_in, n_out, key):

  return dict(weights=jrand.normal(key, (n_in, n_out)) * jnp.sqrt(2/n_in),
                      biases=jnp.ones(shape=(n_out,)))

class BaseJaxNet:
  def __init__(self) -> None:
     self.params = []
     self.loss_fn=None

  @partial(jax.jit, static_argnums=(0,))
  def update(self, lr=0.0001):
    grads = jax.grad(self.loss_fn)(self.params, net, xs, ys)
    return jax.tree_map(lambda p, g: p - lr * g, self.params, grads)
    

class MyJaxNet(BaseJaxNet):
  def __init__(self, in_dim, out_dim,loss_fn) -> None:
    super(MyJaxNet, self).__init__()
    self.key = jrand.PRNGKey(0)
    self.key, subkey = jrand.split(key)
    self.dims = [in_dim,128,128,out_dim]
    self.loss_fn=loss_fn
    for i in range(len(self.dims)-1):
      self.params.append(JaxFC(self.dims[i],self.dims[i+1],subkey))

  def forward(self, params, x):
    *hidden, last = params
    for layer in hidden:
      x = jax.nn.relu(x @ layer['weights'] + layer['biases'])
    return x @ last['weights'] + last['biases']

def loss_fn(params, net, x, y):
  return jnp.mean((net.forward(params, x) - y) ** 2)



def init_mlp_params(layer_widths):
    params = []
    key = jrand.PRNGKey(0)

    for n_in, n_out in zip(layer_widths[:-1], layer_widths[1:]):
        key, subkey = jrand.split(key)
        params.append(dict(weights=jrand.normal(subkey, (n_in, n_out)) * jnp.sqrt(2/n_in),
                       biases=jnp.ones(shape=(n_out,))))
  
    return params

def forward(params, x):
  *hidden, last = params
  for layer in hidden:
    x = jax.nn.relu(x @ layer['weights'] + layer['biases'])
  return x @ last['weights'] + last['biases']

def loss_fn_ori(params, x, y):
  return jnp.mean((forward(params, x) - y) ** 2)

@jax.jit
def update(params, x, y, learning_rate=0.0001):

  grads = jax.grad(loss_fn_ori)(params, x, y)
  return jax.tree_map(
      lambda p, g: p - learning_rate * g, params, grads
  )


layers = [1,128,128,1]
params = init_mlp_params(layers)

key = jrand.PRNGKey(100)
xs = jrand.normal(key, (128,1))
ys = xs ** 2

epoch = 1000
learning_rate = 0.0001
net = MyJaxNet(1,1,loss_fn)

for _ in range(epoch):
  
  print("loss: ", loss_fn(net.params, net, xs, ys))
  #This function works well.
  net.params = update(net.params, xs, ys) 
  #This function only change params once and don't change anymore.
  net.params = net.update()

my class version output:

outside version output:

[adam-hartshorne]

Have you tried using Equinox library ( https://github.com/patrick-kidger/equinox ), it provides a convenient base module from which you can build classes which are correctly structured for JAX.

[tw40210]

Thanks for sharing! It seems pretty good and I don't need to do it from scratch. XD

[jakevdp]

The reason your code is not working as expected is because you are marking self as static when JIT-compiling the update method. This is incompatible with in-place mutation of a class instance, and you are mutating your class instance in the loop when you set or modify the params method.

See https://jax.readthedocs.io/en/latest/faq.html#how-to-use-jit-with-methods for a discussion of this problem and some potential solutions.

[tw40210]

Thanks for answering! I think I know what's broken now.