timecast.learners package¶
Submodules¶
timecast.learners.core module¶
Module contents¶
timecast.learners
-
class
timecast.learners.
AR
(*args, name=None, **kwargs)[source]¶ 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)[source]¶ 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[source]¶ 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)[source]¶ 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)[source]¶ - 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)[source]¶ 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)[source]¶ - 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
[source]¶ 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[source]¶ 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
-
abstract classmethod
-
class
timecast.learners.
Linear
(*args, name=None, **kwargs)[source]¶ 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)[source]¶ 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
[source]¶ 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[source]¶ - 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
-
classmethod
-
class
timecast.learners.
PCR
(*args, name=None, **kwargs)[source]¶ 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)[source]¶ 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[source]¶ 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)[source]¶ 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])[source]¶ - 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)[source]¶ Bases:
flax.nn.base.Module
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
-
classmethod
-
class
timecast.learners.
PredictConstant
(*args, name=None, **kwargs)[source]¶ 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)[source]¶ 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)[source]¶ 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)[source]¶ 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)[source]¶ 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])[source]¶ - 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)[source]¶ 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)[source]¶ 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