fix NaN in fp32 sampling

Author: zdcao121

MODEL_CARD.md

@@ -48,4 +48,4 @@ params, transformer = make_transformer(
 
 ## Speeds, Sizes, Times
 - Both models contain ~13.8M parameters
-- Generating 29,000 crystal samples on a single A100 GPU takes ~1,058 seconds (~37 ms per sample)
+- Generating 45,000 crystal samples on a single A100 GPU takes ~440 seconds (~10 ms per sample)

crystalformer/src/sample.py

@@ -55,6 +55,7 @@ def sample_x(key, h_x, Kx, top_p, temperature, batchsize):
     k = sample_top_p(key_k, x_logit, top_p, temperature)
     loc = loc.reshape(batchsize, Kx)[jnp.arange(batchsize), k]
     kappa = kappa.reshape(batchsize, Kx)[jnp.arange(batchsize), k]
+    kappa = jnp.clip(kappa, a_min=1e-6)  # to avoid numerical issue
     x = sample_von_mises(key_x, loc, kappa/temperature, (batchsize,))
     x = (x+ jnp.pi)/(2.0*jnp.pi) # wrap into [0, 1]
     return key, x 

main.py

@@ -188,7 +188,7 @@
 else:
 
     print("\n========== Start sampling ==========")
-    jax.config.update("jax_enable_x64", True) # to get off compilation warning, and to prevent sample nan lattice 
+    # jax.config.update("jax_enable_x64", True) # to get off compilation warning, and to prevent sample nan lattice 
     #FYI, the error was [Compiling module extracted] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
 
     if args.formula is not None: