Method transform_factorization_spec

axlearn/common/pipeline.py:692–697

(
            spec: Optional[FactorizationSpec],
        )

        specs = VDict(**super().create_parameter_specs_recursively())

        def transform_factorization_spec(
            spec: Optional[FactorizationSpec],
        ) -> Optional[FactorizationSpec]:
            if spec is None:
                return None
            return FactorizationSpec(axes=[None] + list(spec.axes))

        return jax.tree.map(
            lambda spec: dataclasses.replace(