API

Everything below is importable from the top-level anyinit package.

Initializing

anyinit.initialize

initialize(
    model: Any,
    mode: str = "analytic",
    input_spec: Any = None,
    *,
    distribution: str = "normal",
    center: bool = False,
    gains: Mapping[str, float] | None = None,
    seed: int | None = None,
    params: Any = None,
) -> InitReport

Reinitialize model so its activation statistics hold across depth.

Parameters:

Name	Type	Description	Default
model	Any	A torch.nn.Module, a Keras model or layer, or a Flax module.	required
mode	str	"analytic" propagates moments through the graph without running the model; "empirical" measures real batches and assumes nothing.	'analytic'
input_spec	Any	A shape, a batch, or a callable returning a batch. Optional in analytic mode, where it is used only for the validation pass; required in empirical mode and for Flax.	None
distribution	str	Shape of the weight draw: "normal", "uniform" or "sinusoidal".	'normal'
center	bool	Remove each output unit's weight mean, so that a layer discards the mean of its input instead of passing it on.	False
gains	Mapping[str, float] | None	Activation name -> fixed gain, e.g. {"relu": 2**0.5}. Layers feeding that activation get gain / sqrt(fan_in) instead of a solved scale.	None
seed	int | None	Seed for the weight draw. The same seed gives the same weights in every framework.	None
params	Any	Parameter tree for functional frameworks. The new tree comes back as report.params.	None

Returns:

Name	Type	Description
An	InitReport	class:~anyinit.report.InitReport: the scale chosen for every layer, a
	InitReport	depth-stability verdict per activation, and anything AnyInit could not do.

Raises:

Type	Description
ConfigError	An option is invalid. Raised before the model is touched.
BackendNotFoundError	The object is not a model of a supported framework.

anyinit.initialize_params

initialize_params(
    model: Any,
    params: Any,
    input_spec: Any = None,
    mode: str = "analytic",
    **options: Any,
) -> tuple[Any, InitReport]

Functional form for JAX and other immutable-parameter frameworks.

As :func:initialize, but returns the new parameter tree alongside the report::

params, report = anyinit.initialize_params(model, params, (1, 32))

Activations

anyinit.register_activation

register_activation(
    activation: Any = None,
    *,
    name: str | None = None,
    overwrite: bool = False,
) -> Any

Teach AnyInit an activation it does not know.

Usable bare, with arguments, or as a plain call::

@anyinit.register_activation
def relu3(x):
    return torch.relu(x) ** 3

@anyinit.register_activation(name="relu_cubed")
class ReLU3(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x) ** 3

anyinit.register_activation(lambda x: np.maximum(x, 0) ** 3, name="relu3")

The moment map is computed by Gaussian quadrature rather than looked up, and both modes then use it. Registering also makes the activation visible to the tracers, so a module subclass becomes a single graph node instead of being inlined into primitives.

A function written against NumPy works with every backend and needs no framework; one written against a specific framework is evaluated through it. Which it is gets detected.

Parameters:

Name	Type	Description	Default
activation	Any	A callable, or an nn.Module/Keras layer subclass.	None
name	str | None	Registry name. Defaults to the function or class name, lowercased.	None
overwrite	bool	Permit replacing an existing registration of the same name.	False

Returns:

Type	Description
Any	The activation itself, so this works as a decorator.

anyinit.unregister_activation

unregister_activation(name: str) -> None

Remove a registration and discard its cached profile.

anyinit.registered_activations

registered_activations() -> tuple[str, ...]

Names of user-registered activations.

anyinit.gain

gain(name: str, **params: float) -> float

Gain AnyInit gives an activation: weights scaled gain / sqrt(fan_in).

The input standard deviation that lands E[f(z)^2] on one, so gain("relu") is sqrt(2). A bounded activation cannot reach one; its gain is the input scale at the middle of its reachable variance range. Parameters such as negative_slope go in as keywords. Pass the result, or any other value, to initialize(gains=...) to fix it instead of solving for it.

anyinit.activation_profile

activation_profile(
    name: str, backend: Any = None, **params: float
) -> ActivationProfile

Profile of a registered or builtin activation, for inspection.

activation_profile("relu3").chi says whether an activation can hold a signal across depth before anything is built with it. backend is needed only for natively registered code, and is worked out from the registration when omitted.

anyinit.ActivationProfile

Moment map of a single activation, with its stability diagnostics.

kinks property

kinks: list[float]

Points where the activation loses smoothness, found once and reused.

homogeneous_degree property

homogeneous_degree: float | None

Degree p if f(c*x) == c**p * f(x), else None.

A positive answer removes every quadrature call from the hot path, and the degree is the Lyapunov slope of the depth map.

diagnostics property

diagnostics: Diagnostics

Stability verdict for this activation, computed once.

chi property

chi: float

Lyapunov slope d log E[f^2] / d log Var[z] at the operating point.

gain property

gain: float

Standard deviation of the pre-activation this activation is initialized to take.

The classical gain: a layer scaled gain / sqrt(fan_in) and fed a unit second moment hands the activation exactly this. sqrt(2) for ReLU. For a bounded activation, whose E[f^2] = 1 is out of reach, it is the input scale at the middle of its reachable variance range, which is what AnyInit aims at instead.

feasible_m2 property

feasible_m2: tuple[float, float]

Range of second moments the activation can produce.

feasible_var property

feasible_var: tuple[float, float]

Range of variances the activation can produce.

moments

moments(mu: float = 0.0, var: float = 1.0) -> MomentState

Moments of f(z) for z ~ N(mu, var).

mixture_moments

mixture_moments(m_s: float, v_s: float) -> MomentState

Moments when the pre-activation is a Gamma scale mixture of Gaussians.

See :func:quadrature.scale_mixture_moments.

max_moments

max_moments(
    mu: float, var: float, count: float
) -> MomentState

Moments of f(max of count i.i.d. N(mu, var)).

Max pooling after a monotone activation is pooling applied before it, since max_i f(z_i) == f(max_i z_i). Integrating against the pre-activation maximum is therefore exact, where fitting a Gaussian to the activation's own moments is not -- a rectifier's output has half its mass at zero.

local_slope

local_slope(
    mu: float, var: float, *, eps: float = 0.001
) -> float

d log E[f^2] / d log Var[z] at this operating point.

:attr:chi describes the activation at its target; a solver stepping from elsewhere needs the slope where it stands, which for a saturating activation falls toward zero.

solve_input_std

solve_input_std(
    target_m2: float, mu: float = 0.0
) -> float | None

Smallest sigma with E[f(N(mu, sigma^2))^2] == target_m2.

None when the target is outside the activation's reachable range.

The report

anyinit.InitReport dataclass

Result of one initialize call.

objective_error class-attribute instance-attribute

objective_error: float = 0.0

Largest relative gap between an objective and what the solve achieved.

params class-attribute instance-attribute

params: Any = None

New parameter tree, for functional backends such as JAX. None elsewhere.

unscaled class-attribute instance-attribute

unscaled: tuple[tuple[str, str], ...] = ()

Weights left as the framework initialized them, each with the reason.

validation_error class-attribute instance-attribute

validation_error: str | None = None

Why the validation pass that input_spec asked for could not run.

max_deviation property

max_deviation: float | None

Worst prediction-vs-measurement gap beyond sampling noise.

worst_layer property

worst_layer: LayerRecord | None

Layer with the worst unexplained prediction gap.

unstable property

unstable: tuple[StabilityRecord, ...]

Activations that are not depth-stable or whose target is unreachable.

assert_healthy

assert_healthy(tol: float = 0.15) -> None

Raise when the run is not trustworthy. Handy in tests and CI.

to_dict

to_dict() -> dict[str, Any]

Render the report as plain data, for logging.

to_markdown

to_markdown() -> str

Render the report as a Markdown table.

anyinit.LayerRecord dataclass

One scalable layer: what it was given and, if measured, what came out.

constrained class-attribute instance-attribute

constrained: bool = True

Whether the solve had a target here.

A layer feeding a normalization has nothing downstream to aim at, since the normalization discards its scale, so its predicted level is not held against the prediction's accuracy.

fixed class-attribute instance-attribute

fixed: bool = False

Whether the scale came from a user-fixed gain rather than from the solve.

upstream_draws class-attribute instance-attribute

upstream_draws: float = 0.0

Relative variance the weight draws of the layers upstream add to this layer's level.

One draw's deviation from the ensemble compounds with depth, so a layer deep in an unnormalized stack sits further from the prediction than one near the input.

deviation property

deviation: float | None

Relative gap between prediction and measurement, when both exist.

noise_floor property

noise_floor: float | None

Gap size that sampling alone explains.

A measured second moment is an average of squares, so batch sampling contributes about sqrt(2/n). The usually larger term is that an analytic prediction is an ensemble expectation while the model holds one draw, whose layer statistics deviate by about 1/sqrt(units) however much data is pushed through, plus whatever the draws upstream have compounded into its input. Those factors multiply, so the floor is two standard deviations of the logarithm, and is wider above the prediction than below it.

excess_deviation property

excess_deviation: float | None

How far the gap exceeds what sampling explains. Zero means consistent.

anyinit.StabilityRecord dataclass

Depth behavior of one activation.

drift property

drift: float

How far a unit relative error grows over the observed depth.

advice

advice() -> str | None

Plain-language consequence, when the verdict has one.

Backends

anyinit.available_backends

available_backends() -> tuple[str, ...]

Names of the backends whose framework is installed here.

Checked without importing any framework.

anyinit.known_backends

known_backends() -> tuple[str, ...]

Names of every backend AnyInit ships, installed or not.

Errors

Exception hierarchy for AnyInit.

Every failure AnyInit raises derives from :class:AnyInitError, so callers can guard a whole initialize call with one except.

AnyInitError

Bases: Exception

Base class for every AnyInit failure.

BackendNotFoundError

Bases: AnyInitError

No registered backend recognized the object passed as a model.

BackendUnavailableError

Bases: AnyInitError

A backend recognized the model but its framework could not be imported.

TraceError

Bases: AnyInitError

The model's graph could not be captured at full fidelity.

ConfigError

Bases: AnyInitError, ValueError

An argument to initialize is invalid. Raised before the model is touched.