DataProcessor

lafn_inferer.DataProcessor · inherits nnx.Module

__init__(self, scaling_method: str, transform_method: str, in_length: int, out_length: int, in_features: int, out_features: int)

Parameter	Type	Default	Description
scaling_method	str	-	(undocumented)
transform_method	str	-	(undocumented)
in_length	int	-	(undocumented)
out_length	int	-	(undocumented)
in_features	int	-	(undocumented)
out_features	int	-	(undocumented)

pad(data: jax.Array, target_length: Union[int, Sequence[int]], axis: Union[int, Sequence[int]], side: str = "right") -> Tuple[jax.Array, jax.Array]

Pad or trim an N-D array along one or more axes on the specified side ('left' or 'right'). Returns (data, mask) where mask is 1 for original elements and 0 for padding.

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
target_length	Union[int, Sequence[int]]	-	(undocumented)
axis	Union[int, Sequence[int]]	-	(undocumented)
side	str	"right"	(undocumented)

Returns: Tuple[jax.Array, jax.Array] (data, mask)

get_context_index_features(self, context_y, **kwargs) -> jax.Array

Feature engineering based on index.

Parameter	Type	Default	Description
context_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

Returns: jax.Array

get_target_index_features(self, target_y, **kwargs) -> jax.Array

Feature engineering based on index.

Parameter	Type	Default	Description
target_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

Returns: jax.Array

get_index_features(self, context_y, target_y, **kwargs) -> Tuple[jax.Array, jax.Array]

Feature engineering based on index.

Parameter	Type	Default	Description
context_y	-	-	(undocumented)
target_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

Returns: Tuple[jax.Array, jax.Array]

create_time_features(self, context_y, target_y, **kwargs)

Parameter	Type	Default	Description
context_y	-	-	(undocumented)
target_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

scale_context_y(self, context_y, min_val=1e-5, **kwargs)

Parameter	Type	Default	Description
context_y	-	-	(undocumented)
min_val	-	1e-05	(undocumented)
**kwargs	-	-	(undocumented)

scale_target_y(self, target_y, loc, scale, min_val=1e-5, **kwargs)

Parameter	Type	Default	Description
target_y	-	-	(undocumented)
loc	-	-	(undocumented)
scale	-	-	(undocumented)
min_val	-	1e-05	(undocumented)
**kwargs	-	-	(undocumented)

scale_y(self, context_y, target_y, min_val=1e-5, **kwargs)

Parameter	Type	Default	Description
context_y	-	-	(undocumented)
target_y	-	-	(undocumented)
min_val	-	1e-05	(undocumented)
**kwargs	-	-	(undocumented)

scale_data(self, context_x, context_y, target_x, target_y, min_val=1e-5, **kwargs) -> Dict[str, jax.Array]

Parameter	Type	Default	Description
context_x	-	-	(undocumented)
context_y	-	-	(undocumented)
target_x	-	-	(undocumented)
target_y	-	-	(undocumented)
min_val	-	1e-05	(undocumented)
**kwargs	-	-	(undocumented)

Returns: Dict[str, jax.Array]

pad_context_x(self, context_x, **kwargs)

Parameter	Type	Default	Description
context_x	-	-	(undocumented)
**kwargs	-	-	(undocumented)

pad_context_y(self, context_y, **kwargs)

Parameter	Type	Default	Description
context_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

pad_context(self, context_x, context_y, **kwargs)

Parameter	Type	Default	Description
context_x	-	-	(undocumented)
context_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

pad_target_y(self, target_y, **kwargs)

Parameter	Type	Default	Description
target_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

pad_target_x(self, target_x, **kwargs)

Parameter	Type	Default	Description
target_x	-	-	(undocumented)
**kwargs	-	-	(undocumented)

pad_target(self, target_x, target_y, **kwargs)

Parameter	Type	Default	Description
target_x	-	-	(undocumented)
target_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

pad_data(self, context_x, context_y, target_x, target_y, **kwargs) -> Dict[str, jax.Array]

Parameter	Type	Default	Description
context_x	-	-	(undocumented)
context_y	-	-	(undocumented)
target_x	-	-	(undocumented)
target_y	-	-	(undocumented)
**kwargs	-	-	(undocumented)

Returns: Dict[str, jax.Array]

scaler(self, data: jax.Array, loc: jax.Array, scale: jax.Array, min_val: float = 1e-5) -> jax.Array

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
loc	jax.Array	-	(undocumented)
scale	jax.Array	-	(undocumented)
min_val	float	1e-05	(undocumented)

Returns: jax.Array

descaler(self, data: jax.Array, loc: jax.Array, scale: jax.Array, min_val: float = 1e-5) -> jax.Array

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
loc	jax.Array	-	(undocumented)
scale	jax.Array	-	(undocumented)
min_val	float	1e-05	(undocumented)

Returns: jax.Array

z_score_stats(data: jax.Array, axis: int) -> Tuple[jax.Array, jax.Array]

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
axis	int	-	(undocumented)

Returns: Tuple[jax.Array, jax.Array]

minmax_stats(data: jax.Array, axis: int) -> Tuple[jax.Array, jax.Array]

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
axis	int	-	(undocumented)

Returns: Tuple[jax.Array, jax.Array]

robust_stats(data: jax.Array, axis: int) -> Tuple[jax.Array, jax.Array]

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
axis	int	-	(undocumented)

Returns: Tuple[jax.Array, jax.Array]

max_abs_stats(data: jax.Array, axis: int) -> Tuple[jax.Array, jax.Array]

Parameter	Type	Default	Description
data	jax.Array	-	(undocumented)
axis	int	-	(undocumented)

Returns: Tuple[jax.Array, jax.Array]

LAFNInferer

lafn_inferer.LAFNInferer · inherits nnx.Module

__init__(self, model: nnx.Module, data_processor: DataProcessor)

Parameter	Type	Default	Description
model	nnx.Module	-	(undocumented)
data_processor	DataProcessor	-	(undocumented)

__call__(self, context_y: Union[np.ndarray, jax.Array], context_t: Union[np.ndarray, jax.Array, None] = None, target_t: Union[np.ndarray, jax.Array, None] = None, forecast_horizon: Optional[int] = None) -> Tuple[jax.Array, jax.Array, jax.Array]

Forecasting with LAFN model.

Parameter	Type	Default	Description
context_y	Union[np.ndarray, jax.Array]	-	jax.Array of shape (batch_size, context_length, num_y_features) or (context_length, num_y_features) or (context_length,). The historical target values.
context_t	Union[np.ndarray, jax.Array, None]	None	jax.Array of shape (batch_size, context_length, 1) or (context_length, 1) or (context_length,).
target_t	Union[np.ndarray, jax.Array, None]	None	(undocumented)
forecast_horizon	Optional[int]	None	(undocumented)

Returns: Tuple[jax.Array, jax.Array, jax.Array]

get_time_features(self, context_y: Union[np.ndarray, jax.Array], context_t: Union[np.ndarray, jax.Array, None] = None, target_t: Union[np.ndarray, jax.Array, None] = None, forecast_horizon: Optional[int] = None) -> Tuple[jax.Array, jax.Array]

Parameter	Type	Default	Description
context_y	Union[np.ndarray, jax.Array]	-	(undocumented)
context_t	Union[np.ndarray, jax.Array, None]	None	(undocumented)
target_t	Union[np.ndarray, jax.Array, None]	None	(undocumented)
forecast_horizon	Optional[int]	None	(undocumented)

Returns: Tuple[jax.Array, jax.Array]

scale_and_pad_data(self, context_x: jax.Array, context_y: jax.Array, target_x: jax.Array) -> Dict[str, jax.Array]

Parameter	Type	Default	Description
context_x	jax.Array	-	(undocumented)
context_y	jax.Array	-	(undocumented)
target_x	jax.Array	-	(undocumented)

Returns: Dict[str, jax.Array]

forecast(self, context_y: Union[np.ndarray, jax.Array], context_t: Optional[Union[np.ndarray, jax.Array, None]] = None, target_t: Optional[Union[np.ndarray, jax.Array, None]] = None, forecast_horizon: Optional[int] = None)

Parameter	Type	Default	Description
context_y	Union[np.ndarray, jax.Array]	-	(undocumented)
context_t	Optional[Union[np.ndarray, jax.Array, None]]	None	(undocumented)
target_t	Optional[Union[np.ndarray, jax.Array, None]]	None	(undocumented)
forecast_horizon	Optional[int]	None	(undocumented)

pdf(self, y: jax.Array, context_y: Union[np.ndarray, jax.Array], context_t: Optional[Union[np.ndarray, jax.Array, None]] = None, target_t: Optional[Union[np.ndarray, jax.Array, None]] = None, forecast_horizon: Optional[int] = None)

Parameter	Type	Default	Description
y	jax.Array	-	(undocumented)
context_y	Union[np.ndarray, jax.Array]	-	(undocumented)
context_t	Optional[Union[np.ndarray, jax.Array, None]]	None	(undocumented)
target_t	Optional[Union[np.ndarray, jax.Array, None]]	None	(undocumented)
forecast_horizon	Optional[int]	None	(undocumented)

sample(self, context_y: Union[np.ndarray, jax.Array], context_t: Optional[Union[np.ndarray, jax.Array, None]] = None, target_t: Optional[Union[np.ndarray, jax.Array, None]] = None, forecast_horizon: Optional[int] = None, num_samples: int = 100, **kwargs)

Parameter	Type	Default	Description
context_y	Union[np.ndarray, jax.Array]	-	(undocumented)
context_t	Optional[Union[np.ndarray, jax.Array, None]]	None	(undocumented)
target_t	Optional[Union[np.ndarray, jax.Array, None]]	None	(undocumented)
forecast_horizon	Optional[int]	None	(undocumented)
num_samples	int	100	(undocumented)
**kwargs	-	-	(undocumented)
```