timecast.learners package

Submodules

timecast.learners.core module

Module contents

timecast.learners

class timecast.learners.AR(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, timecast.learners.base.FitMixin, flax.nn.base.Module

Note

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Parameters

• x (np.ndarray) – input data

• history_len (int) – length of AR history length

• output_dim (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• history (np.ndarray, optional) – Defaults to None. Optional

• for history (initialization) –

• loc – mean for centering data

• scale – std for normalizing data

Returns

result

Return type

np.ndarray

apply(x: numpy.ndarray, history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Note

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Parameters

• x (np.ndarray) – input data

• history_len (int) – length of AR history length

• output_dim (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• history (np.ndarray, optional) – Defaults to None. Optional

• for history (initialization) –

• loc – mean for centering data

• scale – std for normalizing data

Returns

result

Return type

np.ndarray

classmethod call(x: numpy.ndarray, history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod create(x: numpy.ndarray, history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod create_by_shape(input_specs, x: numpy.ndarray, history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod fit(data: Iterable[Tuple[numpy.ndarray, numpy.ndarray, Any]], input_dim: int, history_len: int, output_dim: int = 1, normalize: bool = True, alpha: float = 1.0, key: jax.numpy.lax_numpy.ndarray = None, history: numpy.ndarray = None, name: str = 'AR', **kwargs) → flax.nn.base.Model

Receives data as an iterable of tuples containing input time series, true time series

Notes

• Use (1, history_len * input_dim) vectors as features (could

consider other representations) * Given a time series of length N and a history of length H, construct N - H + 1 windows * We could infer input_dim from data, but for now, require users to explicitly provide * Assumes we get tuples of time series, not individual time series observations

Parameters

• data – an iterable of tuples containing input/truth pairs of time

• plus any auxiliary value (series) –

• input_dim – number of feature dimensions in input

• history_len – length of history to consider

• output_dim – number of feature dimensions in output

• normalize – zscore data or not

• alpha – for ridge regression

• key – random key for jax random

• history – Any history to pass to AR

• name – name for the top-level module

• kwargs – Extra keyword arguments

Returns

initialized model

Return type

flax.nn.Model

classmethod init(x: numpy.ndarray, history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod init_by_shape(input_specs, x: numpy.ndarray, history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod partial(history_len: int, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

class timecast.learners.ARX(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Parameters

• targets (np.ndarray) – target data

• features (np.ndarray) – feature data

• output_shape (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• history_len (int) – length of history

• constrain – force one parameter per for each slot in history. TODO:

• this better (explain) –

• batched (bool) – first axis is batch axis

Returns

result

Return type

np.ndarray

apply(targets: numpy.ndarray = None, features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Parameters

• targets (np.ndarray) – target data

• features (np.ndarray) – feature data

• output_shape (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• history_len (int) – length of history

• constrain – force one parameter per for each slot in history. TODO:

• this better (explain) –

• batched (bool) – first axis is batch axis

Returns

result

Return type

np.ndarray

classmethod call(targets: numpy.ndarray = None, features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape history_len (int): length of history constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod create(targets: numpy.ndarray = None, features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape history_len (int): length of history constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod create_by_shape(input_specs, targets: numpy.ndarray = None, features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape history_len (int): length of history constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod init(targets: numpy.ndarray = None, features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape history_len (int): length of history constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod init_by_shape(input_specs, targets: numpy.ndarray = None, features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape history_len (int): length of history constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod partial(features: numpy.ndarray = None, history_len: int = 1, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• If batched, assume that first axis is time axis

• If not batched, assume that features and / or targets are one

time step and have no time or batch axis * Delegates much of the error checking to ARXHistory

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape history_len (int): length of history constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

class timecast.learners.ARXHistory(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Parameters

• targets (np.ndarray) – target data

• features (np.ndarray) – feature data

• output_shape (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• constrain – force one parameter per for each slot in history. TODO:

• this better (explain) –

• batched (bool) – first axis is batch axis

Returns

result

Return type

np.ndarray

apply(targets: numpy.ndarray = None, features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Parameters

• targets (np.ndarray) – target data

• features (np.ndarray) – feature data

• output_shape (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• constrain – force one parameter per for each slot in history. TODO:

• this better (explain) –

• batched (bool) – first axis is batch axis

Returns

result

Return type

np.ndarray

classmethod call(targets: numpy.ndarray = None, features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod create(targets: numpy.ndarray = None, features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod create_by_shape(input_specs, targets: numpy.ndarray = None, features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod init(targets: numpy.ndarray = None, features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod init_by_shape(input_specs, targets: numpy.ndarray = None, features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

classmethod partial(features: numpy.ndarray = None, output_shape: Union[Tuple[int, …], int] = 1, constrain: bool = True, batched: bool = False)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Notation

• x = features

• y = targets

• H = history_len

Estimates the following:

hat{y} = sum_{i = 1}^{H + 1} B_i x_{t - i - 1} + a

sum_{i = 1} ^ H A_i y_{t - i} + b

Notes:

• Assumes features and targets have three dimensions: (batch,

history, data). Any extra dimensions are part of the data shape * Doesn’t care if features and targets have different history or input dimensions

Args:

targets (np.ndarray): target data features (np.ndarray): feature data output_shape (Union[Tuple[int, …], int]): int or tuple describing output shape constrain: force one parameter per for each slot in history. TODO: explain this better batched (bool): first axis is batch axis

Returns:

np.ndarray: result

class timecast.learners.FitMixin

Bases: abc.ABC

Mixin class that provides a fit function for offline training / learner initialization

abstract classmethod fit(data: Iterable[Tuple[numpy.ndarray, numpy.ndarray, Any]], input_dim: int, output_dim: int = 1, key: jax.numpy.lax_numpy.ndarray = None, **kwargs) → flax.nn.base.Model

Fit and initialize learner on training data

Notes

• We could infer input_dim from data, but for now, require

users to explicitly provide * output_dim defaults to 1 and is ignored for now

Parameters

• data – an iterable of tuples containing input/truth pairs of time

• plus any auxiliary value (series) –

• input_dim – number of feature dimensions in input

• output_dim – number of feature dimensions in output

• key – random key for jax random

• kwargs – Extra keyword arguments

Returns

initialized model

Return type

flax.nn.Model

class timecast.learners.Linear(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Applies a linear transformation to the inputs along multiple dimensions. :param inputs: The nd-array to be transformed. :param output_shape: tuple of output shape. :param input_axes: tuple with axes to apply the transformation on. :param batch_axes: tuple with batch axes. :param bias: whether to add a bias to the output (default: True). :param dtype: the dtype of the computation (default: float32). :param kernel_init: initializer function for the weight matrix. :param bias_init: initializer function for the bias. :param precision: numerical precision of the computation see jax.lax.Precision

for details.

Returns

The transformed input.

apply(inputs, output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Applies a linear transformation to the inputs along multiple dimensions. :param inputs: The nd-array to be transformed. :param output_shape: tuple of output shape. :param input_axes: tuple with axes to apply the transformation on. :param batch_axes: tuple with batch axes. :param bias: whether to add a bias to the output (default: True). :param dtype: the dtype of the computation (default: float32). :param kernel_init: initializer function for the weight matrix. :param bias_init: initializer function for the bias. :param precision: numerical precision of the computation see jax.lax.Precision

for details.

Returns

The transformed input.

classmethod call(inputs, output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

• inputs – The nd-array to be transformed.

• output_shape – tuple of output shape.

• input_axes – tuple with axes to apply the transformation on.

• batch_axes – tuple with batch axes.

• bias – whether to add a bias to the output (default: True).

• dtype – the dtype of the computation (default: float32).

• kernel_init – initializer function for the weight matrix.

• bias_init – initializer function for the bias.

• precision – numerical precision of the computation see jax.lax.Precision for details.

Returns

The output of the module’s apply function.

Apply docstring:

Applies a linear transformation to the inputs along multiple dimensions.

Returns

The transformed input.

classmethod create(inputs, output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

• inputs – The nd-array to be transformed.

• output_shape – tuple of output shape.

• input_axes – tuple with axes to apply the transformation on.

• batch_axes – tuple with batch axes.

• bias – whether to add a bias to the output (default: True).

• dtype – the dtype of the computation (default: float32).

• kernel_init – initializer function for the weight matrix.

• bias_init – initializer function for the bias.

• precision – numerical precision of the computation see jax.lax.Precision for details.

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Applies a linear transformation to the inputs along multiple dimensions.

Returns

The transformed input.

classmethod create_by_shape(input_specs, inputs, output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

• inputs – The nd-array to be transformed.

• output_shape – tuple of output shape.

• input_axes – tuple with axes to apply the transformation on.

• batch_axes – tuple with batch axes.

• bias – whether to add a bias to the output (default: True).

• dtype – the dtype of the computation (default: float32).

• kernel_init – initializer function for the weight matrix.

• bias_init – initializer function for the bias.

• precision – numerical precision of the computation see jax.lax.Precision for details.

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Applies a linear transformation to the inputs along multiple dimensions.

Returns

The transformed input.

classmethod init(inputs, output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

• inputs – The nd-array to be transformed.

• output_shape – tuple of output shape.

• input_axes – tuple with axes to apply the transformation on.

• batch_axes – tuple with batch axes.

• bias – whether to add a bias to the output (default: True).

• dtype – the dtype of the computation (default: float32).

• kernel_init – initializer function for the weight matrix.

• bias_init – initializer function for the bias.

• precision – numerical precision of the computation see jax.lax.Precision for details.

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Applies a linear transformation to the inputs along multiple dimensions.

Returns

The transformed input.

classmethod init_by_shape(input_specs, inputs, output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Applies a linear transformation to the inputs along multiple dimensions.

Parameters

• inputs – The nd-array to be transformed.

• output_shape – tuple of output shape.

• input_axes – tuple with axes to apply the transformation on.

• batch_axes – tuple with batch axes.

• bias – whether to add a bias to the output (default: True).

• dtype – the dtype of the computation (default: float32).

• kernel_init – initializer function for the weight matrix.

• bias_init – initializer function for the bias.

• precision – numerical precision of the computation see jax.lax.Precision for details.

Returns

The transformed input.

classmethod partial(output_shape, input_axes=-1, batch_axes=(), bias=True, dtype=<class 'jax.numpy.lax_numpy.float32'>, kernel_init=<function variance_scaling.<locals>.init>, bias_init=<function zeros>, precision=None)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

• inputs – The nd-array to be transformed.

• output_shape – tuple of output shape.

• input_axes – tuple with axes to apply the transformation on.

• batch_axes – tuple with batch axes.

• bias – whether to add a bias to the output (default: True).

• dtype – the dtype of the computation (default: float32).

• kernel_init – initializer function for the weight matrix.

• bias_init – initializer function for the bias.

• precision – numerical precision of the computation see jax.lax.Precision for details.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Applies a linear transformation to the inputs along multiple dimensions.

Returns

The transformed input.

class timecast.learners.NewMixin

Bases: object

Mixin class that provides a flax.nn.Model constructor function

classmethod new(shapes: List[Tuple[int, …]], *, key: jax.numpy.lax_numpy.ndarray = None, name: str = None, **kwargs) → flax.nn.base.Model

Parameters

• shapes (List[Tuple[int, ..]]) – shapes for initialization

• key (jnp.ndarray) – key for jax random

• name (str) – identifier for top-level module

• **kwargs – arguments for flax.nn.Module

class timecast.learners.PCR(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, timecast.learners.base.FitMixin, flax.nn.base.Module

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Parameters

• x (np.ndarray) – input data

• history_len (int) – length of AR history length

• output_dim (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• history (np.ndarray, optional) – Defaults to None. Optional

• for history (initialization) –

• loc – mean for centering data

• scale – std for normalizing data

Returns

result

Return type

np.ndarray

apply(x: numpy.ndarray, history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Parameters

• x (np.ndarray) – input data

• history_len (int) – length of AR history length

• output_dim (Union[Tuple[int, ..], int]) – int or tuple

• output shape (describing) –

• history (np.ndarray, optional) – Defaults to None. Optional

• for history (initialization) –

• loc – mean for centering data

• scale – std for normalizing data

Returns

result

Return type

np.ndarray

classmethod call(x: numpy.ndarray, history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Todo:

• AR doesn’t take any history

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod create(x: numpy.ndarray, history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Todo:

• AR doesn’t take any history

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod create_by_shape(input_specs, x: numpy.ndarray, history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Todo:

• AR doesn’t take any history

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod fit(data: Iterable[Tuple[numpy.ndarray, numpy.ndarray, Any]], input_dim: int, history_len: int, output_dim: int = 1, k: int = None, normalize: bool = True, alpha: float = 1.0, key: jax.numpy.lax_numpy.ndarray = None, history: numpy.ndarray = None, name: str = 'PCR', **kwargs) → flax.nn.base.Model

Receives data as an iterable of tuples containing input time series, true time series

Notes

• Use (1, history_len * input_dim) vectors as features (could

consider other representations) * Ignore true value (i.e., look at features only) (should we consider impact of features on dependent variable?) * Given a time series of length N and a history of length H, construct N - H + 1 windows * We could infer input_dim from data, but for now, require users to explicitly provide * Assumes we get tuples of time series, not individual time series observations

Parameters

• data – an iterable of tuples containing input/truth pairs of time

• plus any auxiliary value (series) –

• input_dim – number of feature dimensions in input

• history_len – length of history to consider

• output_dim – number of feature dimensions in output

• k – number of PCA components to keep. Default is min(num_histories,

• normalize – zscore data or not

• input_dim) –

• alpha – Parameter to pass to ridge regression for AR fit

• key – random key for jax random

• history – Any history to pass to AR

• name – name for the top-level module

• kwargs – Extra keyword arguments

Returns

initialized model

Return type

flax.nn.Model

classmethod init(x: numpy.ndarray, history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Todo:

• AR doesn’t take any history

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod init_by_shape(input_specs, x: numpy.ndarray, history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Todo:

• AR doesn’t take any history

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

classmethod partial(history_len: int, projection: numpy.ndarray, output_dim: Union[Tuple[int, …], int] = 1, history: numpy.ndarray = None, loc: Union[numpy.ndarray, float] = None, scale: Union[numpy.ndarray, float] = None)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Todo:

• AR doesn’t take any history

output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape

Note:

• We expect that x is one- or two-dimensional

• We reshape x to ensure its first axis is time and its second axis is input_features

Args:

x (np.ndarray): input data history_len (int): length of AR history length output_dim (Union[Tuple[int, …], int]): int or tuple describing output shape history (np.ndarray, optional): Defaults to None. Optional initialization for history loc: mean for centering data scale: std for normalizing data

Returns:

np.ndarray: result

class timecast.learners.Parallel(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Parameters

• x (np.ndarray) – input data

• modules (List[flax.nn.module]) – list of flax modules

• args (List[dict]) – list of kwargs corresponding to the modules argument to initialize modules

Returns

list of results

Return type

List[np.ndarray]

apply(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Parameters

• x (np.ndarray) – input data

• modules (List[flax.nn.module]) – list of flax modules

• args (List[dict]) – list of kwargs corresponding to the modules argument to initialize modules

Returns

list of results

Return type

List[np.ndarray]

classmethod call(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

List[np.ndarray]: list of results

classmethod create(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

List[np.ndarray]: list of results

classmethod create_by_shape(input_specs, x: numpy.ndarray, learners: List[flax.nn.base.Module])

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

List[np.ndarray]: list of results

classmethod init(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

List[np.ndarray]: list of results

classmethod init_by_shape(input_specs, x: numpy.ndarray, learners: List[flax.nn.base.Module])

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

List[np.ndarray]: list of results

classmethod partial(learners: List[flax.nn.base.Module])

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

List[np.ndarray]: list of results

class timecast.learners.Precomputed(*args, name=None, **kwargs)

Bases: flax.nn.base.Module

Apply function

apply(x, arr)

Apply function

classmethod call(x, arr)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Apply function

classmethod create(x, arr)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Apply function

classmethod create_by_shape(input_specs, x, arr)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Apply function

classmethod init(x, arr)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Apply function

classmethod init_by_shape(input_specs, x, arr)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Apply function

classmethod partial(arr)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Apply function

class timecast.learners.PredictConstant(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Note

• Returns c in the shape of x as the prediction for the next time step

Parameters

• x (np.ndarray) – input data

• c (Real) – prediction

Returns

result

Return type

np.ndarray

Raises

ValueError – if c is not a scalar or does not match the shape of x

apply(x: numpy.ndarray, c: Union[numpy.ndarray, numbers.Real] = 0)

Note

• Returns c in the shape of x as the prediction for the next time step

Parameters

• x (np.ndarray) – input data

• c (Real) – prediction

Returns

result

Return type

np.ndarray

Raises

ValueError – if c is not a scalar or does not match the shape of x

classmethod call(x: numpy.ndarray, c: Union[numpy.ndarray, numbers.Real] = 0)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Note:

• Returns c in the shape of x as the prediction for the next time step

Args:

x (np.ndarray): input data c (Real): prediction

Returns:

np.ndarray: result

Raises:

ValueError: if c is not a scalar or does not match the shape of x

classmethod create(x: numpy.ndarray, c: Union[numpy.ndarray, numbers.Real] = 0)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• Returns c in the shape of x as the prediction for the next time step

Args:

x (np.ndarray): input data c (Real): prediction

Returns:

np.ndarray: result

Raises:

ValueError: if c is not a scalar or does not match the shape of x

classmethod create_by_shape(input_specs, x: numpy.ndarray, c: Union[numpy.ndarray, numbers.Real] = 0)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• Returns c in the shape of x as the prediction for the next time step

Args:

x (np.ndarray): input data c (Real): prediction

Returns:

np.ndarray: result

Raises:

ValueError: if c is not a scalar or does not match the shape of x

classmethod init(x: numpy.ndarray, c: Union[numpy.ndarray, numbers.Real] = 0)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Note:

• Returns c in the shape of x as the prediction for the next time step

Args:

x (np.ndarray): input data c (Real): prediction

Returns:

np.ndarray: result

Raises:

ValueError: if c is not a scalar or does not match the shape of x

classmethod init_by_shape(input_specs, x: numpy.ndarray, c: Union[numpy.ndarray, numbers.Real] = 0)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Note:

• Returns c in the shape of x as the prediction for the next time step

Args:

x (np.ndarray): input data c (Real): prediction

Returns:

np.ndarray: result

Raises:

ValueError: if c is not a scalar or does not match the shape of x

classmethod partial(c: Union[numpy.ndarray, numbers.Real] = 0)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Note:

• Returns c in the shape of x as the prediction for the next time step

Args:

x (np.ndarray): input data c (Real): prediction

Returns:

np.ndarray: result

Raises:

ValueError: if c is not a scalar or does not match the shape of x

class timecast.learners.PredictLast(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Note

• Returns x as the prediction for the next time step

Parameters

x (np.ndarray) – input data

Returns

result

Return type

np.ndarray

apply(x: numpy.ndarray)

Note

• Returns x as the prediction for the next time step

Parameters

x (np.ndarray) – input data

Returns

result

Return type

np.ndarray

classmethod call(x: numpy.ndarray)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Note:

• Returns x as the prediction for the next time step

Args:

x (np.ndarray): input data

Returns:

np.ndarray: result

classmethod create(x: numpy.ndarray)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• Returns x as the prediction for the next time step

Args:

x (np.ndarray): input data

Returns:

np.ndarray: result

classmethod create_by_shape(input_specs, x: numpy.ndarray)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• Returns x as the prediction for the next time step

Args:

x (np.ndarray): input data

Returns:

np.ndarray: result

classmethod init(x: numpy.ndarray)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Note:

• Returns x as the prediction for the next time step

Args:

x (np.ndarray): input data

Returns:

np.ndarray: result

classmethod init_by_shape(input_specs, x: numpy.ndarray)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Note:

• Returns x as the prediction for the next time step

Args:

x (np.ndarray): input data

Returns:

np.ndarray: result

classmethod partial()

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Note:

• Returns x as the prediction for the next time step

Args:

x (np.ndarray): input data

Returns:

np.ndarray: result

class timecast.learners.Sequential(*args, name=None, **kwargs)

Bases: timecast.learners.base.NewMixin, flax.nn.base.Module

Parameters

• x (np.ndarray) – input data

• modules (List[flax.nn.module]) – list of flax modules

• args (List[dict]) – list of kwargs corresponding to the modules argument to initialize modules

Returns

result

Return type

np.ndarray

apply(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Parameters

• x (np.ndarray) – input data

• modules (List[flax.nn.module]) – list of flax modules

• args (List[dict]) – list of kwargs corresponding to the modules argument to initialize modules

Returns

result

Return type

np.ndarray

classmethod call(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

np.ndarray: result

classmethod create(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

np.ndarray: result

classmethod create_by_shape(input_specs, x: numpy.ndarray, learners: List[flax.nn.base.Module])

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

np.ndarray: result

classmethod init(x: numpy.ndarray, learners: List[flax.nn.base.Module])

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

np.ndarray: result

classmethod init_by_shape(input_specs, x: numpy.ndarray, learners: List[flax.nn.base.Module])

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

np.ndarray: result

classmethod partial(learners: List[flax.nn.base.Module])

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Args:

x (np.ndarray): input data modules (List[flax.nn.module]): list of flax modules args (List[dict]): list of kwargs corresponding to the modules

argument to initialize modules

returns:

np.ndarray: result

class timecast.learners.Take(*args, name=None, **kwargs)

Bases: flax.nn.base.Module

Note

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Parameters

• x (np.ndarray) – input data

• index (int) – index to take

Returns

result

Return type

np.ndarray

apply(x: numpy.ndarray, index: int)

Note

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Parameters

• x (np.ndarray) – input data

• index (int) – index to take

Returns

result

Return type

np.ndarray

classmethod call(x: numpy.ndarray, index: int)

Evaluate the module with the given parameters.

Parameters

• params – the parameters of the module. Typically, inital parameter values are constructed using Module.init or Module.init_by_shape.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

The output of the module’s apply function.

Apply docstring:

Note:

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Args:

x (np.ndarray): input data index (int): index to take

Returns:

np.ndarray: result

classmethod create(x: numpy.ndarray, index: int)

Create a module instance by evaluating the model.

DEPRECATION WARNING: create() is deprecated use init() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

Use create_by_shape instead to initialize without doing computation. Initializer functions can depend both on the shape and the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Args:

x (np.ndarray): input data index (int): index to take

Returns:

np.ndarray: result

classmethod create_by_shape(input_specs, x: numpy.ndarray, index: int)

Create a module instance using only shape and dtype information.

DEPRECATION WARNING: create_by_shape() is deprecated use init_by_shape() to initialize parameters and then explicitly create a nn.Model given the module and initialized parameters.

This method will initialize the model without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – other arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and an instance of Model

Apply docstring:

Note:

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Args:

x (np.ndarray): input data index (int): index to take

Returns:

np.ndarray: result

classmethod init(x: numpy.ndarray, index: int)

Initialize the module parameters.

Parameters

• _rng – the random number generator used to initialize parameters.

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Apply docstring:

Note:

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Args:

x (np.ndarray): input data index (int): index to take

Returns:

np.ndarray: result

classmethod init_by_shape(input_specs, x: numpy.ndarray, index: int)

Initialize the module parameters.

This method will initialize the module parameters without computation. Initializer functions can depend on the shape but not the value of inputs.

Parameters

• _rng – the random number generator used to initialize parameters.

• input_specs – an iterable of (shape, dtype) pairs specifying the inputs

• *args – arguments passed to the module’s apply function

• name – name of this module.

• **kwargs – keyword arguments passed to the module’s apply function

Returns

A pair consisting of the model output and the initialized parameters

Example

``` input_shape = (batch_size, image_size, image_size, 3) model_output, initial_params = model.init_by_shape(jax.random.PRNGKey(0),

input_specs=[(input_shape, jnp.float32)])

```

Apply docstring:

Note:

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Args:

x (np.ndarray): input data index (int): index to take

Returns:

np.ndarray: result

classmethod partial(index: int)

Partially applies a module with the given arguments.

Unlike functools.partial this will return a subclass of Module.

Parameters

• name – the name used the module

• **kwargs – the argument to be applied.

Returns

A subclass of Module which partially applies the given keyword arguments.

Apply docstring:

Note:

• Returns x[index] as the prediction for the next time step

• This is a workaround for the case where we have a blackbox series

of predictions (see documentation)

Args:

x (np.ndarray): input data index (int): index to take

Returns:

np.ndarray: result