spark.core.tracers

spark.core.tracers#

Classes#

`BaseTracer`	Base Tracer class
`Tracer`	Multipurpose exponential tracer.
`DoubleTracer`	Multipurpose double exponential tracer.
`SaturableTracer`	Multipurpose saturable exponential tracer.
`SaturableDoubleTracer`	Multipurpose saturable double exponential tracer.
`RUTracer`	Resource-Usage tracer for STP (Short Term Plasticity).

Module Contents#

class spark.core.tracers.BaseTracer(shape, seed=None, dtype=jnp.float16, dt=1.0, **kwargs)[source]#

Bases: flax.nnx.Module, abc.ABC

Base Tracer class

Parameters:

shape (tuple[int, Ellipsis])
seed (int | None)
dtype (Any | None)
dt (float | None)

shape[source]#

rng[source]#

abstractmethod reset()[source]#

Return type:: None

abstractmethod masked_reset(mask)[source]#

Return type:: None

property value: jax.Array[source]#

Abstractmethod:
Return type:: jax.Array

__call__(x, **kwargs)[source]#

Parameters:: x (jax.Array)
Return type:: jax.Array

class spark.core.tracers.Tracer(shape, tau, scale=1, base=0, **kwargs)[source]#

Bases: BaseTracer

Multipurpose exponential tracer.

Parameters:

shape (tuple[int, Ellipsis])
tau (jax.Array | float)
scale (jax.Array | float)
base (jax.Array | float)

scale[source]#

base[source]#

decay[source]#

trace[source]#

reset()[source]#

Return type:: None

masked_reset(mask)[source]#

Return type:: None

property value: jax.Array[source]#

Return type:: jax.Array

class spark.core.tracers.DoubleTracer(shape, tau_1, tau_2, scale_1=1, scale_2=1, base_1=0, base_2=0, *args, **kwargs)[source]#

Bases: BaseTracer

Multipurpose double exponential tracer.

Parameters:

shape (tuple[int, Ellipsis])
tau_1 (jax.Array | float)
tau_2 (jax.Array | float)
scale_1 (jax.Array | float)
scale_2 (jax.Array | float)
base_1 (jax.Array | float)
base_2 (jax.Array | float)

scale_1[source]#

scale_2[source]#

base_1[source]#

base_2[source]#

decay_1[source]#

decay_2[source]#

trace_1[source]#

trace_2[source]#

reset()[source]#

Return type:: None

masked_reset(mask)[source]#

Return type:: None

property value: jax.Array[source]#

Return type:: jax.Array

class spark.core.tracers.SaturableTracer(shape, tau, scale=1, base=0, **kwargs)[source]#

Bases: Tracer

Multipurpose saturable exponential tracer.

Parameters:

shape (tuple[int, Ellipsis])
tau (jax.Array | float)
scale (jax.Array | float)
base (jax.Array | float)

class spark.core.tracers.SaturableDoubleTracer(shape, tau_1, tau_2, scale_1=1, scale_2=1, base_1=0, base_2=0, *args, **kwargs)[source]#