Skip to content

EpochCompressor.fast(config=...) raises TypeError: got multiple values for keyword argument 'config' #37

Description

@jascal

Summary

EpochCompressor.fast(config=…) (and .thorough(config=…)) raise TypeError: got multiple values for keyword argument 'config' even though the docstring states **overrides accepts every constructor kwarg.

Reproducer (polygram 0.1.0)

from polygram import EpochCompressionConfig, EpochCompressor

EpochCompressor.fast(
    sae_checkpoint="sae.safetensors",
    prompts=["hello"],
    layer=8,
    config=EpochCompressionConfig(coverage_target=0.6, max_iterations=2),
)
Traceback (most recent call last):
  ...
  File ".../polygram/compression/epoch.py", line 225, in fast
    return cls._from_preset(defaults, overrides)
  File ".../polygram/compression/epoch.py", line 251, in _from_preset
    return cls(config=preset, **overrides)
TypeError: polygram.compression.epoch.EpochCompressor() got multiple values for keyword argument 'config'

Root cause

polygram/compression/epoch.py:244-251_from_preset passes the preset as config=preset and forwards **overrides without popping a caller-supplied config:

@classmethod
def _from_preset(cls, preset, overrides: dict) -> "EpochCompressor":
    return cls(config=preset, **overrides)

So any caller that passes config= to .fast() / .thorough() collides with the preset. Affects both .fast() (line 212) and .thorough() (line 228).

Expected behavior

The docstring on .fast() says:

**overrides accepts every constructor kwarg, including the required positional inputs (sae_checkpoint, prompts, layer). Tuning kwargs in overrides win over the preset.

Either:

  1. A caller-supplied config= should win (caller has explicitly opted out of the preset), or
  2. The methods should reject config= with a clear error pointing at the constructor (use EpochCompressor(config=…) directly), since .fast() / .thorough() are preset wrappers and accepting a custom config defeats their purpose.

Suggested fix

Option 1 (caller config= wins):

@classmethod
def _from_preset(cls, preset, overrides: dict) -> "EpochCompressor":
    config = overrides.pop("config", preset)
    return cls(config=config, **overrides)

Option 2 (reject collision explicitly):

@classmethod
def _from_preset(cls, preset, overrides: dict) -> "EpochCompressor":
    if "config" in overrides:
        raise TypeError(
            ".fast()/.thorough() supply their own preset config; "
            "for a fully custom config, call EpochCompressor(config=…) directly"
        )
    return cls(config=preset, **overrides)

I'd lean Option 2 — it preserves the named-preset semantics and the error names the right escape hatch.

Downstream impact

sae-forge examples/forge_gpt2_real_sae.py hit this on its end-to-end smoke run. Worked around it on our side by switching to the constructor (EpochCompressor(config=…) instead of EpochCompressor.fast(config=…)) — see jascal/sae-forge@ded117b. Happy to send a polygram PR with whichever option you prefer.

Environment

  • polygram 0.1.0
  • Python 3.11.9, x86_64 macOS (also reproduces on any platform; not Intel-specific)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions