def testGelu(self, approximate):
    def gelu_reference(x):
      return x * scipy.stats.norm.cdf(x)
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng((4, 5, 6), jnp.float32)]
    self._CheckAgainstNumpy(
      gelu_reference, partial(nn.gelu, approximate=approximate), args_maker,
      check_dtypes=False, tol=1e-3 if approximate else None)