spark.nn

Submodules

• spark.nn.brain
• spark.nn.components
• spark.nn.initializers
• spark.nn.interfaces
• spark.nn.neurons

Classes

Module	Base class for Spark Modules
Config	Default class for module configuration.
BaseConfig	Base class for module configuration.
Brain	Abstract brain model.
BrainConfig	Configuration class for Brain's.

Package Contents

class spark.nn.Module(*, config=None, name=None, **kwargs)

Bases: flax.nnx.Module, abc.ABC, Generic[ConfigT, InputT]

Base class for Spark Modules

Parameters:

• config (ConfigT | None)

• name (str | None)

name: str = 'name'

config: ConfigT

default_config: type[ConfigT]

classmethod __init_subclass__(**kwargs)

__built__: bool = False

__allow_cycles__: bool = False

classmethod get_config_spec()

Returns the default configuratio class associated with this module.

Return type:

type[spark.core.config.BaseSparkConfig]

build(input_specs)

Build method.

Parameters:

input_specs (dict[str, spark.core.specs.InputSpec])

Return type:

None

reset()

Reset module to its default state.

set_recurrent_shape_contract(shape=None, output_shapes=None)

Recurrent shape policy pre-defines expected shapes for the output specs.

This is function is a binding contract that allows the modules to accept self connections.

Input:

shape: tuple[int, …], A common shape for all the outputs. output_shapes: dict[str, tuple[int, …]], A specific policy for every single output variable.

NOTE: If both, shape and output_specs, are provided, output_specs takes preference over shape.

Parameters:

• shape (tuple[int, Ellipsis] | None)

• output_shapes (dict[str, tuple[int, Ellipsis]] | None)

Return type:

None

get_recurrent_shape_contract()

Retrieve the recurrent shape policy of the module.

get_input_specs()

Returns a dictionary of the SparkModule’s input port specifications.

Return type:

dict[str, spark.core.specs.InputSpec]

get_output_specs()

Returns a dictionary of the SparkModule’s input port specifications.

Return type:

dict[str, spark.core.specs.OutputSpec]

get_rng_keys(num_keys)

Generates a new collection of random keys for the JAX’s random engine.

Parameters:

num_keys (int)

Return type:

jax.Array | list[jax.Array]

abstractmethod __call__(**kwargs)

Execution method.

Parameters:

kwargs (InputT)

Return type:

ModuleOutput

class spark.nn.Config(**kwargs)

Bases: BaseSparkConfig

Default class for module configuration.

seed: int

dtype: jax.typing.DTypeLike

dt: float

class spark.nn.BaseConfig(**kwargs)

Bases: abc.ABC

Base class for module configuration.

__config_delimiter__: str = '__'

__shared_config_delimiter__: str = '_s_'

__graph_editor_metadata__: dict

classmethod __init_subclass__(**kwargs)

merge(partial={})

Update config with partial overrides.

Parameters:

partial (dict[str, Any])

Return type:

None

diff(other)

Return differences from another config.

Parameters:

other (BaseSparkConfig)

Return type:

dict[str, Any]

validate()

Validates all fields in the configuration class.

Return type:

None

get_field_errors(field_name)

Validates all fields in the configuration class.

Parameters:

field_name (str)

Return type:

list[str]

get_metadata()

Returns all the metadata in the configuration class, indexed by the attribute name.

Return type:

dict[str, Any]

property class_ref: type

Returns the type of the associated Module/Initializer.

NOTE: It is recommended to set the __class_ref__ to the name of the associated module/initializer when defining custom configuration classes. The current class_ref solver is extremely brittle and likely to fail in many different custom scenarios.

Return type:

type

__post_init__()

to_dict()

Serialize config to dictionary

Return type:

dict[str, dict[str, Any]]

classmethod from_dict(dct)

Create config instance from dictionary.

Parameters:

dct (dict)

Return type:

BaseSparkConfig

to_file(file_path)

Export a config instance from a .scfg file.

Parameters:

file_path (str)

Return type:

None

classmethod from_file(file_path)

Create config instance from a .scfg file.

Parameters:

file_path (str)

Return type:

None

__iter__()

Custom iterator to simplify SparkConfig inspection across the entire ecosystem. This iterator excludes private fields.

Output:

field_name: str, field name field_value: tp.Any, field value

Return type:

Iterator[tuple[str, dataclasses.Field, Any]]

get_tree_structure()

Return type:

str

class spark.nn.Brain(config=None, **kwargs)

Bases: spark.core.module.SparkModule

Abstract brain model. This is more a convenience class used to synchronize data more easily.

Parameters:

config (BrainConfig)

config: BrainConfig

resolve_initialization_order()

Resolves the initialization order of the modules.

build(input_specs)

Build method.

Parameters:

input_specs (dict[str, spark.core.specs.InputSpec])

reset()

Resets all the modules to its initial state.

__call__(**inputs)

Update brain’s states.

Parameters:

inputs (spark.core.payloads.SparkPayload)

Return type:

tuple[spark.core.payloads.SparkPayload]

get_spikes_from_cache()

Collect the brain’s spikes.

Return type:

dict

class spark.nn.BrainConfig(**kwargs)

Bases: spark.core.config.BaseSparkConfig

Configuration class for Brain’s.

input_map: dict[str, spark.core.specs.InputSpec]

output_map: dict[str, dict]

modules_map: dict[str, spark.core.specs.ModuleSpecs]

validate()

Validates all fields in the configuration class.

Return type:

None