Fix a few issue in the code.
Add group norm.
Make DP-SGD privacy accountant to be in sync with Tensorflow Privacy.

• Added prototype of ducktyping of Objax variables as JAX arrays
• Added prototype of automatic variable tracing
• Added learning rate scheduler
• Various bugfixes

• Bugfix: removed accidentally introduced TensorFlow dependency.

• Feature: Improved error messages overall
• Feature: Improved BatchNorm numerical stability
• Feature: Objax2Tf for serving objax using TensorFlow
• Feature: New API objax.optimizer.ExponentialMovingAverageModule for easy moving average of a model
• Feature: Automatic broadcasting of scalars for objax.Parallel
• Feature: New optimizer: LARS
• Feature: New API added to functional (lax.scan)
• Feature: Modules can be printed to nicely readable text now (repr)
• Feature: New interpolate API (for images)
• Bugfix: make objax.Sequential work with latest JAX

• Feature: Improved error messages.

• Feature: Extended syntax: allow assigning TrainVar without TrainRef for direction experimentation.

• Feature: Extended padding options or pad and convolution.

• Feature: Modified ResNet_V2 to be Keras compatible.

• Feature: Defaults can be overridden in call for Adam, Momentum.

• BugFix: Layer norm initialization in GPT-2.