strange import behavior of penzai.nn and pz

[hrbigelow]

I'm confused what the intended behavior is between penzai.nn and pz and pz.nn. Here's an example of the confusing behavior. Basically, when you import nn, you don't get everything in that namespace unless you also import pz.

>>> from penzai import nn
>>> help(nn.linear_and_affine.Linear)
Traceback (most recent call last):
  File "<python-input-1>", line 1, in <module>
    help(nn.linear_and_affine.Linear)
         ^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'penzai.nn' has no attribute 'linear_and_affine'
>>> from penzai import pz
>>> help(nn.linear_and_affine.Linear)
Help on class Linear in module penzai.nn.linear_and_affine:

class Linear(penzai.nn.layer.Layer)
 |  Linear(*args, **kwargs)
...

>>> help(pz.nn.linear_and_affine.Linear)
Traceback (most recent call last):
  File "<python-input-4>", line 1, in <module>
    help(pz.nn.linear_and_affine.Linear)
         ^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'penzai.pz.nn' has no attribute 'linear_and_affine'

[danieldjohnson]

Sorry for the confusion! This seems like an artifact of how Python imports work. Right now importing penzai.nn does NOT recursively import all submodules of penzai.nn. However, once somebody imports them (anywhere in the Python program) they get added to the penzai.nn namespace. If you do import penzai.nn.linear_and_affine directly then penzai.nn.linear_and_affine will be available even if you don't use pz.

The pz.nn namespace is intended as a "shortcut" namespace with less nesting, so pz.nn.Linear is an alias for penzai.nn.linear_and_affine.Linear. As a side effect, importing pz.nn also imports the submodules that are being aliased, which causes the behavior you saw.

I'd recommend either:

• using the alias version pz.nn.Linear and not directly accessing things from the nested namespace penzai.nn.linear_and_affine, or
• importing the fully-qualified module penzai.nn.linear_and_affine directly (import penzai.nn.linear_and_affine or from penzai.nn import linear_and_affine and not using the alias namespace pz.nn

Does that make sense?

[hrbigelow]

Yeah, that makes sense. I've been using python for decades now and still get confused over this behavior. I think the confusion is that when I do an import, sometimes I expect that to import everything under it. But of course, whether that happens or not is a choice of the library implementer.

In this case, you have 'from penzai import pz' triggering imports under penzai.nn. That's an odd convention since usually you would expect the designer to automate the imports of things only under the thing being imported.

So, I would say it's better to follow a rule that importing some module only triggers imports of modules under it. But then, I may be confused what "under it" even means, since you can alias things too...