TypeError trying to slice/index ak.Input

[to314as]

When trying to slice a tensor within the ak.AutoModel Framework I get the error "TypeError: 'Input' object is not subscriptable".

I use ak.Input(shape=((None, a, b))) and attempt to slice like x=x[:,i,:] or using a keras.layers.lambda(x: x[:,i,:]) however both versions ended in the same error as described above.
Eg:
input=ak.Input(shape=((None, 20, 2)))
input=keras.layers.lambda(x: x[:,i,:])(input)
model = ak.AutoModel(
inputs=input,
outputs=outputs,
overwrite=True,
max_trials=2,
seed=42,
max_model_size=10**8,
)
model.fit(
train_generator,
validation_data=val_generator,
epochs=100,
)

I observed the same behaviour for ak.StructuredDataInput()
Is it expected behavior that it is not possible to slice a ak.Input() tensor or am I missing a supported way to do this?
If it is not clear from my description: for my model I need to slice/index a tensor after being passed through some layers and then pass the sliced versions through the following versions. Is there any proper way to do this? The only other way to simulate slicing, I could think of, is multiplying with a mask and then summing along the desired dimensions which however performance-wise seems like a suboptimal solution.

Thank you for any advice you can provide!