Axon question: LSTM with Dropout

Questions & Help Questions
,
1

I would like to do something like the following:

input = Axon.input("input", shape: {prediction_days, 1})

model =
  input
  |> Axon.lstm(50)
  |> Axon.dropout(rate: 0.2)
  |> Axon.lstm(50)
  |> Axon.dropout(rate: 0.2)
  |> Axon.lstm(50)
  |> Axon.dropout(rate: 0.2)
  |> Axon.dense(1)

But when I try this out in LiveBook, I get the following error:

** (FunctionClauseError) no function clause matching in Axon.dropout/2    
    
    The following arguments were given to Axon.dropout/2:
    
        # 1
        {#Axon<
           inputs: %{"input" => {60, 1}}
           outputs: "lstm_1_output_sequence"
           nodes: 6
         >,
         {#Axon<
            inputs: %{"input" => {60, 1}}
            outputs: "lstm_1_c_hidden_state"
            nodes: 6
          >,
          #Axon<
            inputs: %{"input" => {60, 1}}
            outputs: "lstm_1_h_hidden_state"
            nodes: 6
          >}}
    
        # 2
        [rate: 0.2]
    
    Attempted function clauses (showing 1 out of 1):
    
        def dropout(%Axon{} = x, opts)

The Python code I am trying to replicate looks like this:

model = Sequential()

model.add(LSTM(units=50, return_sequences=True, input_shape=[x_train.shape[1],  1]))
model.add(Dropout(0.2))
model.add(LSTM(units=50, return_sequences=True))
model.add(Dropout(0.2))
model.add(LSTM(units=50))
model.add(Dropout(0.2))
model.add(Dense(units=1))

model.compile(optimizer='adam', loss='mean_squared_error')
model.fit(x_train, y_train, epochs=25, batch_size=32)

Any help? Thanks in advance!