Improved nn.py

I changed the nn definitions to use jnp.clip (hopefully faster) and changed it so that user specified beta values are single learnable constants for the layer - might change again in the future to allow flexible specification of either scalar or vectors for beta.