NickleDave

import mlflow

# There are two ways to create parent/child runs in MLflow.

# (1) The most common way is to use the fluent
# mlflow.start_run API, passing nested=True:
with mlflow.start_run():
  num_trials = 10
  mlflow.log_param("num_trials", num_trials)
  best_loss = 1e100

NickleDave

"""
Simplified Implementation of the Linear Recurrent Unit
------------------------------------------------------
We present here a simplified JAX implementation of the Linear Recurrent Unit (LRU). 
The state of the LRU is driven by the input $(u_k)_{k=1}^L$ of sequence length $L$ 
according to the following formula (and efficiently parallelized using an associative scan): 
$x_{k} = \Lambda x_{k-1} +\exp(\gamma^{\log})\odot (B u_{k})$, 
and the output is computed at each timestamp $k$ as follows: $y_k = C x_k + D u_k$. 
In our code, $B,C$ follow Glorot initialization, with $B$ scaled additionally by a factor 2 
to account for halving the state variance by taking the real part of the output projection.