spark.core.module

Attributes

ConfigT
InputT

Classes

ModuleOutput	Spark module output template
SparkMeta	Metaclass for Spark Modules.
SparkModule	Base class for Spark Modules

Module Contents

spark.core.module.ConfigT

spark.core.module.InputT

class spark.core.module.ModuleOutput

Bases: TypedDict

Spark module output template

Initialize self. See help(type(self)) for accurate signature.

class spark.core.module.SparkMeta

Bases: flax.nnx.module.ModuleMeta

Metaclass for Spark Modules.

class spark.core.module.SparkModule(*, 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 configuration 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.PortSpecs])

Return type:

None

reset()

Reset module to its default state.

set_contract_output_specs(contract_specs)

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:

contract_specs: dict[str, PortSpecs], A dictionary with a contract for the output specs.

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

Parameters:

contract_specs (dict[str, spark.core.specs.PortSpecs])

Return type:

None

get_contract_output_specs()

Retrieve the recurrent spec policy of the module.

Return type:

dict[str, spark.core.specs.PortSpecs] | None

get_input_specs()

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

Return type:

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

get_output_specs()

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

Return type:

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

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

__repr__()

inspect()

Returns a formated string of the datastructure.

Return type:

str

checkpoint(path, overwrite=False)

Return type:

None

classmethod from_checkpoint(path, safe=True)

Return type:

SparkModule