Help getting an Optimizer working

[harryjulian]

Hi, I've got a bug in my code that I'm struggling to fix or understand. I've compared my code to the examples and couldn't find a fix, so I thought I'd ask here!

I'm using Optax to train a Layer of a Forward-Forward Neural Net.

Here's some of the code:

LayerParams = list[Array, Array]
Layer = list[LayerParams, GradientTransformation]

def train_layer(key: KeyArray, X: Array, y: Array, layer: Layer, learning_rate: float, epochs: int, theta: int):

  def train_step(stepkey, params, X_pos, X_neg, opt_state):
    stepkey, subkey = jax.random.split(stepkey, 2)
    loss_val, grads, _, _ = activations(params[0], params[1], X_pos, X_neg, theta)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    weights, biases = apply_updates(params, updates)
    return subkey, loss_val, params, opt_state

  # Unpack & Initialise
  params, optimizer = layer[0], layer[1]
  opt_state = optimizer.init(params)

  # Optimize
  for i in range(epochs):

    # Split Key, get Positive and Negative Data
    key, subkey = jax.random.split(key, 2)
    X_pos = overlay_and_flatten(X, y)
    y_neg = swap_imperfect(subkey, y)
    X_neg = overlay_and_flatten(X, y_neg)

    key, loss_val, params, opt_state = train_step(
      subkey, 
      params,
      X_pos,
      X_neg,
      opt_state
    )
    if i % 100 == 0: print(f'Epoch {i}, loss: {loss_val}')

  return layer

I keep running into this error:

File /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/optax/_src/transform.py:84, in update_moment.<locals>.<lambda>(g, t)
     [81](file:///Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/optax/_src/transform.py?line=80) def update_moment(updates, moments, decay, order):
     [82](file:///Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/optax/_src/transform.py?line=81)   """Compute the exponential moving average of the `order`-th moment."""
     [83](file:///Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/optax/_src/transform.py?line=82)   return jax.tree_util.tree_map(
---> [84](file:///Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/optax/_src/transform.py?line=83)       lambda g, t: (1 - decay) * (g ** order) + decay * t, updates, moments)

TypeError: can't multiply sequence by non-int of type 'float'

On this line

updates, opt_state = optimizer.update(grads, opt_state, params)

Any ideas would be super valuable:)