[Chex] Assertion assert_equal_shape failed: Arrays have different shapes: [(32, 2), (32,)].

[mwitiderrick]

My X_train is in the shape (32, 60, 1) while y_train is in the shape
(32,) yet I have this shapes error. What am I doing wrong?

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Input In [32], in <cell line: 28>()
     26         print(f"Epoch: {epoch}, train loss: {train_metrics['loss']}, train accuracy: {train_metrics['accuracy'] * 100}")
     27     return train_state
---> 28 trained_model_state = train_model()

Input In [32], in train_model()
     18 def train_model():
     19     for epoch in range(1, num_epochs + 1):
---> 20         train_state, train_metrics = train_one_epoch(state)
     21         training_loss.append(train_metrics['loss'])
     22         training_accuracy.append(train_metrics['accuracy'])

Input In [31], in train_one_epoch(state)
      3 batch_metrics = []
      4 for x, y in get_train_batches():
----> 5     state, metrics = train_step(state, x, y)
      6     batch_metrics.append(metrics)
      8 batch_metrics_np = jax.device_get(batch_metrics)  

    [... skipping hidden 14 frame]

Input In [29], in train_step(state, x, y)
      6   return loss, logits
      7 grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
----> 8 (_, logits), grads = grad_fn(state.params)
      9 state = state.apply_gradients(grads=grads)
     10 metrics = compute_metrics(logits, y)

    [... skipping hidden 8 frame]

Input In [29], in train_step.<locals>.loss_fn(params)
      3 def loss_fn(params):
      4   logits = LSTMModel().apply({'params': params}, x)
----> 5   loss = jnp.mean(optax.l2_loss(logits, y))
      6   return loss, logits

File ~/miniforge3/lib/python3.9/site-packages/optax/_src/loss.py:55, in l2_loss(predictions, targets)
     52 chex.assert_type([predictions], float)
     53 if targets is not None:
     54   # Avoid broadcasting logic for "-" operator.
---> 55   chex.assert_equal_shape((predictions, targets))
     56 errors = (predictions - targets) if (targets is not None) else predictions
     57 return 0.5 * (errors)**2

File ~/miniforge3/lib/python3.9/site-packages/chex/_src/asserts_internal.py:197, in chex_assertion.<locals>._chex_assert_fn(*args, **kwargs)
    195 else:
    196   try:
--> 197     host_assertion(*args, **kwargs)
    198   except jax.errors.ConcretizationTypeError as exc:
    199     msg = ("Chex assertion detected `ConcretizationTypeError`: it is very "
    200            "likely that it tried to access tensors' values during tracing. "
    201            "Make sure that you defined a jittable version of this Chex "
    202            "assertion.")

File ~/miniforge3/lib/python3.9/site-packages/chex/_src/asserts_internal.py:157, in make_static_assertion.<locals>._static_assert(custom_message, custom_message_format_vars, include_default_message, exception_type, *args, **kwargs)
    154     custom_message = custom_message.format(*custom_message_format_vars)
    155   error_msg = f"{error_msg} [{custom_message}]"
--> 157 raise exception_type(error_msg)

AssertionError: [Chex] Assertion assert_equal_shape failed: Arrays have different shapes: [(32, 2), (32,)].

[gerdm]

It's hard to say without your model definition, but I suspect is because logits is a 2-dimensional array, whereas your target variable is 1-dimensional. Have you tried setting yhat = logits.argmax(axis=-1) and then setting loss = jnp.mean(optax.l2_loss(yhat, y))?

[mwitiderrick]

@gerdm thanks for taking the time to respond. I tried this and got a different error message

AssertionError: [Chex] Assertion assert_type failed: Error in type compatibility check: input 0 has type int32 but expected <class 'float'>.

Here is my model definition

from flax import linen as nn

class LSTMModel(nn.Module):
    def setup(self):
        lstm_layer = nn.scan(nn.OptimizedLSTMCell,
                               variable_broadcast="params",
                               split_rngs={"params": False},
                               in_axes=1, 
                               out_axes=1,
                               # length=max_len,
                               reverse=False)
        self.lstm1 = lstm_layer()
        self.dense1 = nn.Dense(256)
        self.lstm2 = lstm_layer()
        self.dense2 = nn.Dense(128)
        self.lstm3 = lstm_layer()
        self.dense3 = nn.Dense(64)
        self.dense4 = nn.Dense(2)
        
    # @nn.remat    
    def __call__(self, x_batch):
        
        carry, hidden = nn.OptimizedLSTMCell.initialize_carry(jax.random.PRNGKey(0), batch_dims=(len(x_batch),), size=128)
        (carry, hidden), x = self.lstm1((carry, hidden), x_batch)
        
        x = self.dense1(x)
        x = nn.relu(x)
        
        carry, hidden = nn.OptimizedLSTMCell.initialize_carry(jax.random.PRNGKey(0), batch_dims=(len(x_batch),), size=64)
        (carry, hidden), x = self.lstm2((carry, hidden), x)
        
        x = self.dense2(x)
        x = nn.relu(x)
        
        carry, hidden = nn.OptimizedLSTMCell.initialize_carry(jax.random.PRNGKey(0), batch_dims=(len(x_batch),), size=32)
        (carry, hidden), x = self.lstm3((carry, hidden), x)
        
       
        x = self.dense3(x)
        x = nn.relu(x)
        x = self.dense4(x[:, -1])
        return x

When I convert the logits to float, I get a different error

AssertionError: [Chex] Assertion assert_equal_shape failed: Arrays have different shapes: [(7,), (32,)].

[mwitiderrick]

@cgarciae any idea how to fix this?

[cgarciae]

@mwitiderrick can you put some minimum code to reproduce this so I can run it?

[mwitiderrick]

@cgarciae thanks for your response. I am attaching a link to a Colab 🙏 https://colab.research.google.com/drive/1SJDwBMIRzBEmmN9dcxs0f1agU2ltH4CN?usp=sharing