Question on reshape

[roger.lee]

Would like to understand the difference between the two, in the compute_attention_output in the assignment

x = jnp.reshape(x, (-1, n_heads, seqlen, d_head))

x = jnp.reshape(x, (x.shape[0] / n_heads, n_heads, seqlen, d_head))

The first one produce no error while the second one incurs error. Is there a difference in terms of implementation? Thank you!

[arvyzukai]

Hi @roger.lee

There is a bug in a test case:

C4_W2_Assignment UNQ_C4 unit test has a bug: the testing data is wrong

According to the doc string, the input x shape is (n_batch X n_heads, seqlen, d_head) So when reshaping x shape (n_batch, n_heads, seqlen, d_head), I use n_batch = x.shape[0] // n_heads x = jnp.reshape(x, (n_batch, n_heads, seqlen, d_head)) However, the unit test failed with a stack trace. If I use x = jnp.reshape(x, (-1, n_heads, seqlen, d_head)) Then it passes. Looking into the failed test case, I found: x.shape = (6, 2, 3) However, n_heads = 3, seqlen = 2, d_head = 2. So n_batch sho…

I think you are talking about this error (test unit case error and not the code error)?

Cheers