Reduces the ops in the graph by
1) casting stochastic discretization `step_size` to `value_type`, rather than casting `value` to float32.
2) applying run length encoding layer-wise to each tensor, rather than concatenating structures to a single tensor for encoding.

PiperOrigin-RevId: 491695106