Mechanize cannot work with Dropout?

[neel04]

I'm using equinox and wanted to use mechanize. However, in the jax.lax.cond @ this line apparently the PyTree structures used in the cond are different thus triggering the error:

Code

self.schedule_fn = optax.warmup_cosine_decay_schedule(init_value=self.lr / 2, peak_value=self.lr,
                                                      warmup_steps=self.warmup_steps, decay_steps=total_steps)

# optimizer with weight decay
optim = optax.chain(
    optax.lion(learning_rate=self.schedule_fn, weight_decay=self.weight_decay),
    optax.clip_by_global_norm(self.grad_clip),
    optax.apply_every(self.accum_steps)
)

optim = optax.contrib.mechanize(optim)
opt_state = optim.init(eqx.filter(model, eqx.is_array))

...

updates, opt_state = optim.update(grads, opt_state, model) # ERRORS OUT

Error

Traceback (most recent call last):
  File "/Users/neel/Documents/research/ReAct_Jax/train_model.py", line 104, in <module>
    main(key)
  File "/Users/neel/Documents/research/ReAct_Jax/train_model.py", line 72, in main
    trainer.train()
  File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 246, in train
    loss, model, opt_state = make_step(model, seq, label, pad_mask, rndm_n, rndm_k,
  File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py", line 206, in __call__
    return self._call(False, args, kwargs)
  File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_module.py", line 875, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
  File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/equinox/_jit.py", line 200, in _call
    out = self._cached(dynamic_donate, dynamic_nodonate, static)
  File "/Users/neel/Documents/research/ReAct_Jax/ReAct/utils/trainer.py", line 68, in make_step
    updates, opt_state = optim.update(grads, opt_state, model)
  File "/Users/neel/miniconda3/envs/react_jax/lib/python3.10/site-packages/optax/contrib/mechanic.py", line 139, in update_fn
    x0 = jax.lax.cond(state.count == 0, lambda: params, lambda: state.x0)
TypeError: true_fun and false_fun output must have same type structure, got PyTreeDef(CustomNode(React[('pos_enc', 'embed_layer', 'main_block', 'post_ln', 'out_head'), ........

# Truncated, but it prints out both the PyTrees

Triaging the problem

I put the 2 PyTrees whose structure is dissimilar and computed the diff. I think the first PyTree on the left is params and the one on the right is state.x0.

https://www.diffchecker.com/sYMTXWWi/

It appears that using Dropout is the common denominator in each difference. This is the first training step, so perhaps it marks some fields as None before they're updated? I'm not sure.

Note: This is the definition of eqx.nn.Dropout as in the equinox package. The core logic is really simple.

Is MECHANIC not supposed to work with Dropout? I'm open to PR-ing fixes if that's of any help.

[neel04]

@fabianp @adefazio Any ideas?

[fabianp]

Thanks for reporting this. Looping also @nullstring on this issue.

Is this issue specific to mechanize ? I.e., if you replace optax.optim.mechanize with (say) optax.sgd then the code runs fine?

[neel04]

Yep, it works fine then