Fix Flax usage comments (#1211)

Commit

3 years ago

Fix Flax usage comments (#1211) * Fix Flax usage comments (they didn't work). * Spell out dtype * make style