Raphaël Boige

bug: can't type flax.struct.dataclass with vmapped functions

3

Hey, runtime type-checking seems to fail when providing a Flax dataclass to a vmapped function. I wasn't able to find related resources . Here is a minimal reproduction with the...

The Sudoku problem has many symmetries and permutations to which it is invariant or equivariant. The currently used generators works by sampling existing instances from a fixed database. To improve...

Pretraining loss used in the code

2

Hello and thanks for making all this code public! I have a question regarding the pretraining loss actually used in the training, I understand through looking at the pretraining code...