How to Discretization on ragged 3d?

How to use layers.Discretization() across dimention when a ragged tensor innermost dimension is 1 ? Is there any layers.Reshape() trick?
Like [[[10],[20],[30]]] to become [[[1],[2],[3]]]
And [[[10],[20],[30],[40]]] to become [[[1],[2],[3],[4]]]
Since it is ragged, reshape() need to flexiable right?