add experimental support for mutable array

[cgarciae]

Examples

jit

@struct.dataclass
class Params:
  w: nx.Variable
  b: nx.Variable
  count: nx.Variable

  @classmethod
  def create(cls, din, dout: int):
    return cls(
      w=nx.Variable(jnp.zeros((din, dout))),
      b=nx.Variable(jnp.zeros((dout,))),
      count=nx.Variable(0),
    )

params = Params.create(3, 4)

@jax.jit # <=== 🎉
def linear(params, x):
  params.count[...] += 1
  return x @ params.w + params.b[None]

x = jnp.ones((1, 3))
y = linear(params, x)

assert y.shape == (1, 4)
assert params.count[...] == 1

scan

class Count(nx.Variable):
  pass

@struct.dataclass
class Weights:
  w: nx.Param
  b: nx.Param
  count: Count

  @staticmethod
  def create(num_blocks: int, dim: int):
    return Weights(
      w=nx.Param(jnp.zeros((num_blocks, dim, dim))),
      b=nx.Param(jnp.zeros((num_blocks, dim))),
      count=Count(0),
    )

weights = Weights.create(3, 5)
treedef, params, counts = nx.split(weights, nx.Param, Count)

def forward(x: jax.Array, params):
  weights = nx.merge(treedef, params, counts)
  x = x @ weights.w + weights.b[None]
  weights.count[...] += 1
  return x, None

x = jnp.ones((1, 5))
y, _ = jax.lax.scan(forward, x, params)

assert y.shape == (1, 5)
assert weights.count[...] == 3

training loop

class Linear(nx.Pytree):
  __nodes__ = ('w', 'b', 'count')

  def __init__(self, din: int, dout: int, *, rngs: nx.Rngs):
    self.w = nx.Param(jax.random.uniform(rngs(), (din, dout), jnp.float32))
    self.b = nx.Param(jnp.zeros((dout,), jnp.float32))
    self.count = nx.Variable(0)

  def __call__(self, x: jax.Array):
    self.count[...] += 1
    return x @ self.w + self.b[None]

class MLP(nx.Pytree):
  __nodes__ = ('linear1', 'linear2')

  def __init__(self, din: int, dmid: int, dout: int, *, rngs: nx.Rngs):
    self.linear1 = Linear(din, dmid, rngs=rngs)
    self.linear2 = Linear(dmid, dout, rngs=rngs)

  def __call__(self, x: jax.Array):
    return self.linear2(jax.nn.relu(self.linear1(x)))

model = MLP(1, 64, 1, rngs=nx.Rngs(0))

@jax.jit
def train_step(model, x, y):
  treedef, params, rest = nx.split(model, nx.Param, ...)

  def loss_fn(params):
    model = nx.merge(treedef, params, rest)
    return jnp.mean((model(x) - y) ** 2)

  loss, grads = jax.value_and_grad(loss_fn)(nx.freeze(params))

  def sgd(p: nx.Variable, g: jax.Array):
    p[...] -= 0.1 * g

  jax.tree.map(sgd, params, grads)
  return loss

loss = 100
for _ in range(10):
  x = np.linspace(0, 1, 32)[:, None]
  y = 0.8 * x**2 + 0.1 + np.random.normal(0, 0.1, size=x.shape)

  loss = train_step(model, x, y)

assert loss < 1.0