added flax.typing

[chiamp]

Added flax/typing.py file to unify typing annotations.

Few things to note:

• Filter is defined differently in NNX, so I did not unify NNX's definition under the flax.typing definition
• Axis is defined differently in NNX, so I did not unify NNX's definition under the flax.typing definition
• In addition to regular JAX initializers, Partitioned objects can also be used as kernel_init and bias_init arguments. Currently to accommodate this, the Initializer type annotation is Union[jax.nn.initializers.Initializer, Callable[..., Any]]. The more correct annotation would be Union[jax.nn.initializers.Initializer, Callable[..., Partitioned[Any]]], but then the linter will complain about code that tries to call Array methods or access Array attributes like .reshape and .shape on the initialized kernel or bias since Partitioned[Any] does not have any of these methods or attributes. The linter will also complain about functions that take in the initialized kernel or bias as an input argument, because they expect an Array and not Partitioned[Any].

[codecov-commenter]

Codecov Report

Attention: 7 lines in your changes are missing coverage. Please review.

Comparison is base (74bdc7b) 56.19% compared to head (96d2d5f) 56.12%.
Report is 2 commits behind head on main.

Files	Patch %	Lines
flax/training/dynamic_scale.py	0.00%	3 Missing ⚠️
flax/linen/partitioning.py	0.00%	2 Missing ⚠️
flax/cursor.py	0.00%	1 Missing ⚠️
flax/linen/linear.py	92.30%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3624      +/-   ##
==========================================
- Coverage   56.19%   56.12%   -0.08%     
==========================================
  Files         101      102       +1     
  Lines       12214    12180      -34     
==========================================
- Hits         6864     6836      -28     
+ Misses       5350     5344       -6     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.