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.

[danieldjohnson]

Also, mind running uv run pyink penzai to re-format the changed files using pyink so that the checks pass?

Thanks much!

[amifalk]

By the way, the .vscode setup files didn't work for me, and I had to manually call pyink with uvx pyink penzai. Edit: For some reason uv didn't automatically install the dev dependencies, but after calling uv pip install .[dev] it worked.

[danieldjohnson]

Looks like pytype is raising an error here. I don't think it knows how to refine the type of self.weights. Would you mind either:

• refactoring this so that you first assign to a local variable weights = self.weights and then change the rest of the function to refer to weights,
• or just adding a comment # pytype: disable=attribute-error here to tell pytype you know what you're doing?

[danieldjohnson]

(You should be able to run the typechecker yourself with uv run pytype --jobs auto penzai as long as the dev dependencies are installed.)

[amifalk]

Hm, for someone reason locally disabling the error doesn't seem to work

    if isinstance(
        self.weights,
        Parameter | ParameterValue,
    ) and (  # pytype: disable=attribute-error
        self.weights.label.endswith(".weights")
    ):
      error_prefix = (  # pytype: disable=attribute-error
          f"({self.weights.label[:-8]}) "
      )
    else:
      error_prefix = ""
``

and neither does assigning self.weights to weights. Globally ignoring the attribute error does seem to work though.

[danieldjohnson]

Hm that's weird. Does it work if you add

# pytype: disable=attribute-error

on a line on its own before the if block, and

# pytype: enable=attribute-error

on a line after?

[amifalk]

That did the trick!

[danieldjohnson]

By the way, the .vscode setup files didn't work for me, and I had to manually call pyink with uvx pyink penzai. Edit: For some reason uv didn't automatically install the dev dependencies, but after calling uv pip install .[dev] it worked.

I'm not sure if this is the same thing you're running into, but I've been running into an issue where uvx (as run through VSCode) sometimes selects a default Python version that Black/Pyink doesn't support. I've tried to fix it in #97.