More informative structure mismatch error message in Linear layer

[amifalk]

As mentioned in #89. With the following changes,

from penzai import pz
import jax

linear = pz.nn.Linear.from_config("MyModel/MyComplexLayer/Linear_2", 
                                  jax.random.PRNGKey(0),
                                  input_axes={"feature": 3},
                                  output_axes={})

linear(pz.nx.ones({"feature": 4}))

will throw an error with information about the layer name:

penzai.core.shapecheck.StructureMismatchError: (MyModel/MyComplexLayer/Linear_2) Mismatch while checking structures:
At root: Named shape mismatch between value {'feature': 4} and pattern {**var('B'), 'feature': 3}:
  Axis 'feature': Actual size 4 does not match expected 3

This makes it much easier to debug shape errors in complex models with many layers.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.