Use new `jax.ops.index_mul` in dropout instead of `jax.ops.index_update`. Per my benchmarks it seems to make Dropout much faster and practically of negligible cost. For example batch size = 16 with pooling, 20 Dropout layers cost + 0.4 seconds vs + 80 seconds for the old implementation. I wonder if this is too good to be true, and there's a mistake somewhere?...
Use new `jax.ops.index_mul` in dropout instead of `jax.ops.index_update`. Per my benchmarks it seems to make Dropout much faster and practically of negligible cost. For example batch size = 16 with pooling, 20 Dropout layers cost + 0.4 seconds vs + 80 seconds for the old implementation. I wonder if this is too good to be true, and there's a mistake somewhere?...

PiperOrigin-RevId: 306816000

google/neural-tangents