Objax

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