Skip to content

Make FactorizedTensor.__init_subclass__'s name argument optional - #38

Closed
speedyshreya wants to merge 1 commit into
tensorly:mainfrom
speedyshreya:fix-optional-name-kwarg
Closed

speedyshreya wants to merge 1 commit into
tensorly:mainfrom
speedyshreya:fix-optional-name-kwarg

Conversation

@speedyshreya

Copy link
Copy Markdown

Fixes tensorly/tensorly#614 (issue filed on the main tensorly repo; the affected code belongs in this repo).

Problem
torch.nn.utils.parametrize.register_parametrization dynamically builds a subclass of the module it's applied to, via type(f"Parametrized{cls.name}", (cls,), dct). Because init_subclass declared name with no default, this raised TypeError: init_subclass() missing 1 required positional argument: 'name' — PyTorch has no knowledge of the convention and no way to satisfy it. This affects any FactorizedTensor subclass; the reporter hit it via ComplexTuckerTensor inside a neuraloperator FNO model.

Fix
Default name to ''. The name == '' branch already handles subclasses that shouldn't be registered (it's how TensorizedTensor opts out), so this simply makes that path reachable for callers that can't pass the argument at all.

I deliberately did not default to cls.name as suggested in the issue, because that would set cls._name = 'ParametrizedTuckerTensor' — corrupting the public .name property, documented as returning the factorization name — and add a junk parametrizedtuckertensor key to _factorizations on every parametrization. Inheriting the parent's _name keeps .name == 'Tucker', which is correct.

Explicitly-named subclasses are unaffected, and the existing warning for hand-written subclasses omitting a name still fires.

Tests
Added a parametrized regression test covering Tucker, CP, Dense and their Complex variants. TT is excluded as it has no direct top-level parameter to target.

Note: 3 pre-existing failures in test_trl.py (TypeError: init() got an unexpected keyword argument 'verbose') are unrelated — they reproduce on main without this change.

@speedyshreya

Copy link
Copy Markdown
Author

closing this

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant