added Einsum layer

[chiamp]

Added Einsum layer. See API docs here.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 35.00000% with 26 lines in your changes are missing coverage. Please review.

Project coverage is 58.56%. Comparing base (c83ee86) to head (8e56442).
Report is 27 commits behind head on main.

Files	Patch %	Lines
flax/linen/linear.py	35.00%	26 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3710      +/-   ##
==========================================
- Coverage   59.10%   58.56%   -0.55%     
==========================================
  Files         103      103              
  Lines       12426    12419       -7     
==========================================
- Hits         7345     7273      -72     
- Misses       5081     5146      +65     

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

[cgarciae]

use canonical variables name

Suggested change

	>>> params = layer.init(jax.random.key(0), jnp.ones((3, 4, 5)))
	>>> jax.tree_map(jnp.shape, params)
	>>> variables = layer.init(jax.random.key(0), jnp.ones((3, 4, 5)))
	>>> jax.tree_map(jnp.shape, variables)