def train_SBL(model: DeepMoD, data: torch.Tensor, target: torch.Tensor, optimizer, extra_params, sparsity_scheduler, split=0.8, exp_ID: str=None, log_dir: str=None, max_iterations: int=10000, write_iterations: int=25, **convergence_kwargs) -> None:
    'Trains the DeepMoD model. This function automatically splits the data set in a train and test set. \n\n    Args:\n        model (DeepMoD):  A DeepMoD object.\n        data (torch.Tensor):  Tensor of shape (n_samples x (n_spatial + 1)) containing the coordinates, first column should be the time coordinate.\n        target (torch.Tensor): Tensor of shape (n_samples x n_features) containing the target data.\n        optimizer ([type]):  Pytorch optimizer.\n        sparsity_scheduler ([type]):  Decides when to update the sparsity mask.\n        split (float, optional):  Fraction of the train set, by default 0.8.\n        exp_ID (str, optional): Unique ID to identify tensorboard file. Not used if log_dir is given, see pytorch documentation.\n        log_dir (str, optional): Directory where tensorboard file is written, by default None.\n        max_iterations (int, optional): [description]. Max number of epochs , by default 10000.\n        write_iterations (int, optional): [description]. Sets how often data is written to tensorboard and checks train loss , by default 25.\n    '
    logger = Logger(exp_ID, log_dir)
    sparsity_scheduler.path = logger.log_dir
    (t, a, l) = extra_params
    n_train = int((split * data.shape[0]))
    n_test = (data.shape[0] - n_train)
    (data_train, data_test) = torch.split(data, [n_train, n_test], dim=0)
    (target_train, target_test) = torch.split(target, [n_train, n_test], dim=0)
    M = 12
    N = data_train.shape[0]
    threshold = 10000.0
    convergence = Convergence(**convergence_kwargs)
    for iteration in torch.arange(0, max_iterations):
        (prediction, time_derivs, thetas) = model(data_train)
        tau_ = torch.exp(t)
        alpha_ = torch.min(torch.exp(a), torch.tensor(100000000.0, dtype=torch.float32))
        lambda_ = torch.min(torch.exp(l), torch.tensor(20000.0, dtype=torch.float32))
        y = time_derivs[0]
        X = (thetas[0] / torch.norm(thetas[0], dim=0, keepdim=True))
        p_MSE = ((N / 2) * (((tau_ * torch.mean(((prediction - target_train) ** 2), dim=0)) - t) + np.log((2 * np.pi))))
        A = (torch.diag(lambda_) + ((alpha_ * X.T) @ X))
        mn = ((lambda_ < threshold)[:, None] * (((alpha_ * torch.inverse(A)) @ X.T) @ y))
        E = ((alpha_ * torch.sum(((y - (X @ mn)) ** 2))) + ((mn.T @ torch.diag(lambda_)) @ mn))
        p_reg = ((1 / 2) * (((E + torch.sum(torch.log(torch.diag(A)[(lambda_ < threshold)]))) - (torch.sum(l[(lambda_ < threshold)]) + (N * a))) - (N * np.log((2 * np.pi)))))
        MSE = torch.mean(((prediction - target_train) ** 2), dim=0)
        Reg = torch.stack([torch.mean(((dt - (theta @ coeff_vector)) ** 2)) for (dt, theta, coeff_vector) in zip(time_derivs, thetas, model.constraint_coeffs(scaled=False, sparse=True))])
        loss = torch.sum((p_MSE + p_reg))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if ((iteration % write_iterations) == 0):
            with torch.no_grad():
                prediction_test = model.func_approx(data_test)[0]
                MSE_test = torch.mean(((prediction_test - target_test) ** 2), dim=0)
            _ = model.sparse_estimator(thetas, time_derivs)
            logger(iteration, loss, MSE, Reg, model.constraint_coeffs(sparse=True, scaled=True), model.constraint_coeffs(sparse=True, scaled=False), model.estimator_coeffs(), MSE_test=MSE_test, p_MSE=p_MSE, p_reg=p_reg, tau=tau_, alpha=alpha_, lambda_=lambda_, mn=mn)
            update_sparsity = sparsity_scheduler(iteration, torch.sum(MSE_test), model, optimizer)
            if update_sparsity:
                model.constraint.sparsity_masks = model.sparse_estimator(thetas, time_derivs)
            l1_norm = torch.sum(torch.abs(torch.cat(model.constraint_coeffs(sparse=True, scaled=True), dim=1)))
            converged = convergence(iteration, l1_norm)
            if converged:
                break
    logger.close(model)
Trains the DeepMoD model. This function automatically splits the data set in a train and test set. 

Args:
    model (DeepMoD):  A DeepMoD object.
    data (torch.Tensor):  Tensor of shape (n_samples x (n_spatial + 1)) containing the coordinates, first column should be the time coordinate.
    target (torch.Tensor): Tensor of shape (n_samples x n_features) containing the target data.
    optimizer ([type]):  Pytorch optimizer.
    sparsity_scheduler ([type]):  Decides when to update the sparsity mask.
    split (float, optional):  Fraction of the train set, by default 0.8.
    exp_ID (str, optional): Unique ID to identify tensorboard file. Not used if log_dir is given, see pytorch documentation.
    log_dir (str, optional): Directory where tensorboard file is written, by default None.
    max_iterations (int, optional): [description]. Max number of epochs , by default 10000.
    write_iterations (int, optional): [description]. Sets how often data is written to tensorboard and checks train loss , by default 25.
def train_SBL(model: DeepMoD, data: torch.Tensor, target: torch.Tensor, optimizer, extra_params, sparsity_scheduler, split=0.8, exp_ID: str=None, log_dir: str=None, max_iterations: int=10000, write_iterations: int=25, **convergence_kwargs) -> None:
    'Trains the DeepMoD model. This function automatically splits the data set in a train and test set. \n\n    Args:\n        model (DeepMoD):  A DeepMoD object.\n        data (torch.Tensor):  Tensor of shape (n_samples x (n_spatial + 1)) containing the coordinates, first column should be the time coordinate.\n        target (torch.Tensor): Tensor of shape (n_samples x n_features) containing the target data.\n        optimizer ([type]):  Pytorch optimizer.\n        sparsity_scheduler ([type]):  Decides when to update the sparsity mask.\n        split (float, optional):  Fraction of the train set, by default 0.8.\n        exp_ID (str, optional): Unique ID to identify tensorboard file. Not used if log_dir is given, see pytorch documentation.\n        log_dir (str, optional): Directory where tensorboard file is written, by default None.\n        max_iterations (int, optional): [description]. Max number of epochs , by default 10000.\n        write_iterations (int, optional): [description]. Sets how often data is written to tensorboard and checks train loss , by default 25.\n    '
    logger = Logger(exp_ID, log_dir)
    sparsity_scheduler.path = logger.log_dir
    (t, a, l) = extra_params
    n_train = int((split * data.shape[0]))
    n_test = (data.shape[0] - n_train)
    (data_train, data_test) = torch.split(data, [n_train, n_test], dim=0)
    (target_train, target_test) = torch.split(target, [n_train, n_test], dim=0)
    M = 12
    N = data_train.shape[0]
    threshold = 10000.0
    convergence = Convergence(**convergence_kwargs)
    for iteration in torch.arange(0, max_iterations):
        (prediction, time_derivs, thetas) = model(data_train)
        tau_ = torch.exp(t)
        alpha_ = torch.min(torch.exp(a), torch.tensor(100000000.0, dtype=torch.float32))
        lambda_ = torch.min(torch.exp(l), torch.tensor(20000.0, dtype=torch.float32))
        y = time_derivs[0]
        X = (thetas[0] / torch.norm(thetas[0], dim=0, keepdim=True))
        p_MSE = ((N / 2) * (((tau_ * torch.mean(((prediction - target_train) ** 2), dim=0)) - t) + np.log((2 * np.pi))))
        A = (torch.diag(lambda_) + ((alpha_ * X.T) @ X))
        mn = ((lambda_ < threshold)[:, None] * (((alpha_ * torch.inverse(A)) @ X.T) @ y))
        E = ((alpha_ * torch.sum(((y - (X @ mn)) ** 2))) + ((mn.T @ torch.diag(lambda_)) @ mn))
        p_reg = ((1 / 2) * (((E + torch.sum(torch.log(torch.diag(A)[(lambda_ < threshold)]))) - (torch.sum(l[(lambda_ < threshold)]) + (N * a))) - (N * np.log((2 * np.pi)))))
        MSE = torch.mean(((prediction - target_train) ** 2), dim=0)
        Reg = torch.stack([torch.mean(((dt - (theta @ coeff_vector)) ** 2)) for (dt, theta, coeff_vector) in zip(time_derivs, thetas, model.constraint_coeffs(scaled=False, sparse=True))])
        loss = torch.sum((p_MSE + p_reg))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if ((iteration % write_iterations) == 0):
            with torch.no_grad():
                prediction_test = model.func_approx(data_test)[0]
                MSE_test = torch.mean(((prediction_test - target_test) ** 2), dim=0)
            _ = model.sparse_estimator(thetas, time_derivs)
            logger(iteration, loss, MSE, Reg, model.constraint_coeffs(sparse=True, scaled=True), model.constraint_coeffs(sparse=True, scaled=False), model.estimator_coeffs(), MSE_test=MSE_test, p_MSE=p_MSE, p_reg=p_reg, tau=tau_, alpha=alpha_, lambda_=lambda_, mn=mn)
            update_sparsity = sparsity_scheduler(iteration, torch.sum(MSE_test), model, optimizer)
            if update_sparsity:
                model.constraint.sparsity_masks = model.sparse_estimator(thetas, time_derivs)
            l1_norm = torch.sum(torch.abs(torch.cat(model.constraint_coeffs(sparse=True, scaled=True), dim=1)))
            converged = convergence(iteration, l1_norm)
            if converged:
                break
    logger.close(model)