eztaox.kernels.eqx_utils

Utility methods for Equinox modules.

Functions

find_param_by_name(→ list | None)

Find a leaf parameter in an Equinox module by name.

Module Contents

find_param_by_name(module: equinox.Module, name: str) → list | None

Find a leaf parameter in an Equinox module by name.

Parameters:

• module (eqx.Module) – The Equinox module to search in.

• name (str) – The name of the parameter to find.

Returns:

The parameter if found, None otherwise.

Return type:

list | None