vmap seems incompatible with CNN (haiku)

[Xin-yang-Liu]

I am trying to vmap a 2D CNN, which is defined by using haiku. Every time jax will raise a warning says

None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.

followed by

Conv: (f32[80,16,16,16]{3,2,1,0}, u8[0]{0}) custom-call(f32[80,32,22,22]{3,2,1,0}, f32[16,32,7,7]{3,2,1,0}, f32[16]{0}), window={size=7x7}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convBiasActivationForward", backend_config={"conv_result_scale":1,"activation_mode":"kRelu","side_input_scale":0,"leakyrelu_alpha":0}

When I disable the vmap the warning will disappear. I know I could stack every other axis to the batch dimension to avoid using vmap, however, it will be much convenient for me to keep the axises. I wonder if this is a bug in Jax or Haiku.

Thanks in advance!

[hawkinsp]

That sounds like an XLA bug: it had trouble compiling a particular convolution.

If this reproduces with an up to date jax and jaxlib, please file a bug.

[Xin-yang-Liu]

It turns out this warning message is actually caused by OOM...