def linear_scheduler_with_warmup(
    learning_rate: float,
    num_inputs_train: int,
    train_total_batch_size: int,
    num_train_epochs: int,
    num_warmup_steps: int,
) -> Callable[[int], jnp.array]:
    """
    Create a linear scheduler with a warmup

    Parameters
    ----------
    learning_rate: float
        Learning rate value
    num_inputs_train: int
        Number of input data points in the training set.
    train_total_batch_size: int
        Total training batch size.
    num_train_epochs: int
        Number of training epochs.
    num_warmup_steps: int
        Number of warmup steps. At the end of the warm-up phase the learning rate value will be `learning_rate`.

    Returns
    -------
    Callable[[int], jnp.array]
        The scheduler function
    """
    steps_per_epoch = num_inputs_train // train_total_batch_size
    num_train_steps = steps_per_epoch * num_train_epochs
    warmup_fn = optax.linear_schedule(
        init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps
    )
    decay_fn = optax.linear_schedule(
        init_value=learning_rate,
        end_value=0,
        transition_steps=num_train_steps - num_warmup_steps,
    )
    schedule_fn = optax.join_schedules(
        schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps]
    )
    return schedule_fn

    Create a linear scheduler with a warmup

    Parameters
    ----------
    learning_rate: float
        Learning rate value
    num_inputs_train: int
        Number of input data points in the training set.
    train_total_batch_size: int
        Total training batch size.
    num_train_epochs: int
        Number of training epochs.
    num_warmup_steps: int
        Number of warmup steps. At the end of the warm-up phase the learning rate value will be `learning_rate`.

    Returns
    -------
    Callable[[int], jnp.array]
        The scheduler function
    
fortuna/utils/optimizer.py