Delays

rex.base.DelayDistribution

A delay distribution data structure.

reset(rng: jax.Array) -> DelayDistribution

Reset the distribution (e.g. random number generator).

Parameters:

• rng (Array) –

  random number generator

Returns:

• DelayDistribution ( DelayDistribution ) –

  the reset distribution

sample(shape: Union[int, Tuple] = None) -> Tuple[DelayDistribution, jax.Array]

Sample from the distribution.

Parameters:

• shape (Union[int, Tuple], default: None ) –

  the shape of the sample

Returns:

• Tuple[DelayDistribution, Array] –

  The new distribution and the sample

quantile(q: float) -> float

Compute the quantile of the distribution.

Parameters:

• q (float) –

  the quantile value

Returns:

• float –

  The quantile value

mean() -> float

Compute the mean of the distribution.

Returns:

• float –

  The mean value

pdf(x: float) -> float

Compute the probability density function of the distribution.

Parameters:

• x (float) –

  the value at which to compute the pdf

Returns:

• float –

  The pdf value

---

rex.base.StaticDist

Bases: DelayDistribution

A wrapper around distrax distributions to make them compatible with the DelayDistribution interface.

Attributes:

• rng (Array) –

  the random number generator

• dist (Distribution) –

  the (static) distrax distribution (e.g. Normal, MixtureSameFamily, etc.)

create(dist: distrax.Distribution) -> StaticDist classmethod

Create a static distribution.

Parameters:

• dist (Distribution) –

  the distrax distribution

Returns:

• StaticDist –

  The static distribution

reset(rng: jax.Array) -> StaticDist

Reset the random number generator.

Parameters:

• rng (Array) –

  random number generator

Returns:

• DelayDistribution ( StaticDist ) –

  the reset distribution

sample(shape: Union[int, Tuple] = None) -> Tuple[StaticDist, jax.Array]

Sample from the distribution.

Parameters:

• shape (Union[int, Tuple], default: None ) –

  the shape of the sample

Returns:

• Tuple[StaticDist, Array] –

  The new distribution and the sample

quantile(q: float) -> Union[float, jax.typing.ArrayLike]

Compute the quantile of the distribution.

Parameters:

• q (float) –

  the quantile value

Returns:

• Union[float, ArrayLike] –

  The quantile value

mean() -> float

Compute the mean of the distribution.

Returns:

• float –

  The mean value

pdf(x: float) -> float

Compute the probability density function of the distribution.

Parameters:

• x (float) –

  the value at which to compute the pdf

Returns:

• float –

  The pdf value

---

rex.base.TrainableDist

Bases: DelayDistribution

Attributes:

• alpha (Union[float, ArrayLike]) –

  the value between [0, 1]

• min (Union[float, ArrayLike]) –

  the minimum expected delay

• max (Union[float, ArrayLike]) –

  the maximum expected delay

• interp (str) –

  the interpolation method ("zoh", "linear", "linear_real_only"). "zoh": zero-order hold interpolation between received messages. "linear": linear interpolation between received messages. "linear_real_only": Only start interpolating between received messages.

create(delay: Union[float, jax.typing.ArrayLike], min: Union[float, jax.typing.ArrayLike], max: Union[float, jax.typing.ArrayLike], interp: str = 'zoh') -> TrainableDist classmethod

Creates a trainable distribution. Converts the delay to alpha, which is the value between [0, 1].

Parameters:

• delay (Union[float, ArrayLike]) –

  the desired delay

• min (Union[float, ArrayLike]) –

  the minimum expected delay

• max (Union[float, ArrayLike]) –

  the maximum expected delay

• interp (str, default: 'zoh' ) –

  the interpolation method ("zoh", "linear", "linear_real_only").

Returns: The trainable distribution

reset(rng: jax.Array) -> TrainableDist

Does nothing as the distribution is deterministic.

Parameters:

• rng (Array) –

  random number generator

Returns:

• DelayDistribution ( TrainableDist ) –

  the reset distribution

sample(shape: Union[int, Tuple] = None) -> Tuple[TrainableDist, jax.Array]

Sample from the distribution.

Parameters:

• shape (Union[int, Tuple], default: None ) –

  the shape of the sample

Returns:

• Tuple[TrainableDist, Array] –

  The new distribution and the sample

quantile(q: float) -> float

Compute the quantile of the distribution.

As the distribution is deterministic, the quantile is trivially calculated as the constant value of the distribution.

Parameters:

• q (float) –

  the quantile value

Returns:

• float –

  The quantile value

mean() -> Union[float, jax.typing.ArrayLike]

Compute the mean of the distribution.

Returns:

• Union[float, ArrayLike] –

  The mean value

pdf(x: float) -> Union[jax.Array, float]

Compute the probability density function of the distribution.

Parameters:

• x (float) –

  the value at which to compute the pdf

Returns:

• Union[Array, float] –

  The pdf value

equivalent(other: DelayDistribution) -> bool

Check if two delay distributions are equivalent

Parameters:

• other (DelayDistribution) –

  the other delay distribution

Returns:

• bool –

  True if the two delay distributions are equivalent, False otherwise

window(rate_out: Union[float, int]) -> int

Compute the additional window size needed for the delay distribution.

Parameters:

• rate_out (Union[float, int]) –

  the output rate of the connection

Returns:

• int –

  The additional window size needed for the delay distribution

apply_delay(rate_out: float, input: InputState, ts_start: Union[float, jax.typing.ArrayLike]) -> InputState

Apply the delay to the input state.

The delay is determined by the delay distribution of the connection.

Parameters:

• rate_out (float) –

  the output rate of the connection

• input (InputState) –

  the input state

• ts_start (Union[float, ArrayLike]) –

  the start time of the computation

Returns:

• InputState –

  The input state with the delay applied. This reduces the window size of the input by self.window(rate_out).

get_alpha(delay: Union[float, jax.typing.ArrayLike]) -> Union[float, jax.typing.ArrayLike]

Get the alpha value of the delay distribution.

Parameters:

• delay (Union[float, ArrayLike]) –

  the delay value

Returns:

• Union[float, ArrayLike] –

  The alpha value