Correct initializer of the model's params with input of different sizes

[TheGlitch2701]

Hi everyone,
my question is the following: is it possible to initialize a flax model with a dummy input that can take into account different input sizes?

For instance what I'm doing now is the following:

variables = module.init(rng, jnp.ones([1,A,3]))

However, when I have to do my training process I have as input some clouds of points that have the shape of (1,B,3) with B is different not only from A (the 2nd dimension of the dummy input) but also each input has different B.

For instance: B_1 != B_2 != B_3 != ... !=A .

So in the end: is there a way to tell my model that the input it will recieve is different from the dummy one?

Thank you in advance guys!