Shape Inference + Module Hooks (compact)

[cgarciae]

Changes

• Adds the tx.next_key() function and the tx.rng_key() context manager.
• Module.init new has the following behavior:
  • Accepts an optional inputs argument and runs the forward method if given.
  • Set the given key in the context so tx.next_key() can be used.
  • Accepts a call_method: str which defines the method to call, "__call__" used by default.
• Modules will now be initialized if constructed within @tx.compact functions when called by init.
• Adds @tx.compact_module decorator that can turn any function into a Module with a compact __call__ as the decorated function.
• New Crossentropy loss that generalizes BinaryCrossentropy, CategoricalCrossentropy and SparseCategoricalCrossentropy.
• New Flatten layers.