danielmamay / grokking Goto Github PK

Thanks for your implementation. However, I believe there is a bug at the following lines:

https://github.com/danielmamay/grokking/blob/a1b9d96a92cfae3c10aa0b3b76b66b4cea041635/grokking/data.py#L36C1-L37C1

ALL_OPERATIONS[operation](x, y, p) is a lambda function. Under default setting: lambda x, y, p: (x*y % p, y, x).

This lambda function returns a tuple with three elements in the following order: solution-of-the-equation, y-value, x-value

The solution-of-the-equation should be the label. However, in the code, we have:

x, y, z = ALL_OPERATIONS[operation](x, y, p)
results = z.remainder(p)
inputs = torch.stack([x, op, y, eq], dim=1)
labels = results

Where the x is treated as the input, but the value of x is actually the solution to x*y%p.

Can you please conduct an investigation on this matter?

Also, I do not quite understand why you have z.remainder(p) here (assuming z is passed by the solution to the equation).

Grokking is clearly present at 25% training data, though!