Struggling with objax, Modules, StateVar, and python classes

[logan-dunbar]

I am having trouble understanding how to use objax Modules with python classes, and manually updating StateVars. In the code below, basically I want to have a container holding num Test objects, each of which has their own train_op and Adam optimizer. It's a silly example and not sure the maths makes sense, but it seems to highlight the issue.

If I run the code with num = 1 below, it Jit compiles and runs fine. As soon as I set num = 2 I get an UnexpectedTracerError which points to the opt_model(0.1, dE) line. The documentation seems to say that StateVars are used for manually tracking variables, in my case I have a set of manually updated variables and Adam updated variables that are updated using the same loss function. What am I doing wrong here? (The functools.partial seems to break things even with num = 1).

import objax
import numpy as np
import jax.numpy as jnp

class Test(objax.Module):
  def __init__(self, x, y):
    self.x = objax.TrainVar(jnp.array([x]))
    self.y = objax.StateVar(jnp.array(y.copy()))

  # @functools.partial(jax.jit, static_argnums=(0,))
  def energy(self, s, t):
    self.y.assign(self.x.value * s + t * self.y.value)
    return jnp.sum(self.y.value)

class Container:
  def __init__(self, num=2):
    self.objs = {}
    self.ops = {}

    for i in range(num):
      self.objs[i] = Test(0.1, np.random.randn(5))

      opt_model = objax.optimizer.Adam(self.objs[i].vars())
      energy = objax.GradValues(self.objs[i].energy, self.objs[i].vars())

      @objax.Function.with_vars(self.objs[i].vars() + opt_model.vars())
      def train_op(_s, _t):
        dE, E = energy(_s, _t)
        opt_model(0.1, dE)

      self.ops[i] = objax.Jit(train_op)

  def update(self):
    for i, obj in self.objs.items():
      self.ops[i](np.random.random(), np.random.random())

def leak_issue():
  num = 2
  container = Container(num=num)
  # with jax.checking_leaks():
  container.update()

if __name__ == '__main__':
  leak_issue()

[logan-dunbar]

Weirdly, I can replace the energy function with return jnp.sum(self.x.value * s) which doesn't use the StateVar at all, and then even more weirdly it runs through the first of train_op fine, but then fails on the second call. Somehow the definition of the function, or the closure on def train_op or something is causing the values to leak across object instance boundaries. It fails in the Adam update call where the train_var already arrives as a Traced object for some reason, where on the first call it doesn't... :/

[logan-dunbar]

Even weirder, I've tracked it down to line 67 in objax/optimizer/adam.py:

p.value -= lr * m.value * functional.rsqrt(v.value + self.eps)

With num = 3 , on the first run, if I stop the debugger just before running that line, only the one being updated is Traced, both the others are just normal Arrays as expected, but after stepping over that line, just for the first train_op update, then the first and the last have both changed to Traced<>, but not the second. This leads me to believe there might be some sort of caching going on here somewhere?

Before step:

After step:

[logan-dunbar]

Definitely some caching/closure issue, if I create a module level function like so:

def get_train_op(model):
  opt_model = objax.optimizer.Adam(model.vars())
  energy = objax.GradValues(model.energy, model.vars())

  def train_op(_s, _t):
    dE, E = energy(_s, _t)
    opt_model(0.1, dE)

  return objax.Jit(train_op, model.vars() + opt_model.vars())

and then use it to create my train_op:

self.ops[i] = get_train_op(self.objs[i])

then everything works as expected, I can use the StateVar like normal as well... Happy to have gotten to the bottom of it, but if someone could please explain to me what is going on?

I literally traced the moment it changes the wrong TrainVar to a Traced<> is when the p.value = calls the TrainVar.assign function and returns, that is when it switches. I'm not able to find why or how the wrong variable is referenced, or where the code is that switches it out under the hood. Some illumination would be much appreciated :)