cleverhans
cleverhans copied to clipboard
Update fast_gradient_method.py
The lastest Stax in JAX should be jax.example_libraries.stax as documented here: https://jax.readthedocs.io/en/latest/jax.example_libraries.stax.html