Lowering of normalization call when using flax

[pratnali-aws]

Hello All,
This is a newbie question:
I have the following simple model that uses BatchNormalization:

import jax
import jax.numpy as jnp
from flax import linen as nn
import optax
class SimpleCNN(nn.Module):
    @nn.compact
    def __call__(self, x, training):
        x1 = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x2 = nn.BatchNorm(use_running_average=not training)(x1)
        x3 = nn.relu(x2)
        return x3

key = jax.random.PRNGKey(0)
input_shape = (1, 64, 64, 3)
dummy_input = jnp.ones(input_shape)
model = SimpleCNN()
params = model.init(key, dummy_input, True)

@jax.jit
def train_step(params, x):
    predictions = model.apply(params, x, False)
    return predictions
    
xla_comp = jax.xla_computation(train_step)(params, dummy_input)
hlo_text = xla_comp.as_hlo_text()

and it is lowered to

HloModule xla_computation_train_step, entry_computation_layout={(f32[32]{0}, f32[32]{0}, f32[32]{0}, f32[32]{0}, f32[32]{0}, /*index=5*/f32[3,3,3,32]{3,2,1,0}, f32[1,64,64,3]{3,2,1,0})->(f32[1,64,64,32]{3,2,1,0})}

relu.8 {
  Arg_0.9 = f32[1,64,64,32]{3,2,1,0} parameter(0)
  constant.10 = f32[] constant(0)
  broadcast.11 = f32[1,64,64,32]{3,2,1,0} broadcast(constant.10), dimensions={}
  ROOT maximum.12 = f32[1,64,64,32]{3,2,1,0} maximum(Arg_0.9, broadcast.11)
}

train_step_0.13 {
  Arg_6.20 = f32[1,64,64,3]{3,2,1,0} parameter(6)
  Arg_5.19 = f32[3,3,3,32]{3,2,1,0} parameter(5)
  convolution.23 = f32[1,64,64,32]{3,2,1,0} convolution(Arg_6.20, Arg_5.19), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f
  Arg_4.18 = f32[32]{0} parameter(4)
  reshape.24 = f32[1,1,1,32]{3,2,1,0} reshape(Arg_4.18)
  broadcast.25 = f32[1,1,1,32]{3,2,1,0} broadcast(reshape.24), dimensions={0,1,2,3}
  reshape.26 = f32[1,32]{1,0} reshape(broadcast.25)
  broadcast.27 = f32[1,64,64,32]{3,2,1,0} broadcast(reshape.26), dimensions={0,3}
  add.28 = f32[1,64,64,32]{3,2,1,0} add(convolution.23, broadcast.27)
  Arg_0.14 = f32[32]{0} parameter(0)
  reshape.29 = f32[1,1,1,32]{3,2,1,0} reshape(Arg_0.14)
  broadcast.31 = f32[1,1,1,32]{3,2,1,0} broadcast(reshape.29), dimensions={0,1,2,3}
  reshape.32 = f32[1,32]{1,0} reshape(broadcast.31)
  broadcast.33 = f32[1,64,64,32]{3,2,1,0} broadcast(reshape.32), dimensions={0,3}
  subtract.34 = f32[1,64,64,32]{3,2,1,0} subtract(add.28, broadcast.33)
  Arg_1.15 = f32[32]{0} parameter(1)
  reshape.30 = f32[1,1,1,32]{3,2,1,0} reshape(Arg_1.15)
  constant.21 = f32[] constant(1e-05)
  broadcast.22 = f32[1,1,1,32]{3,2,1,0} broadcast(constant.21), dimensions={}
  add.35 = f32[1,1,1,32]{3,2,1,0} add(reshape.30, broadcast.22)
  rsqrt.36 = f32[1,1,1,32]{3,2,1,0} rsqrt(add.35)
  Arg_3.17 = f32[32]{0} parameter(3)
  reshape.37 = f32[1,1,1,32]{3,2,1,0} reshape(Arg_3.17)
  multiply.38 = f32[1,1,1,32]{3,2,1,0} multiply(rsqrt.36, reshape.37)
  broadcast.39 = f32[1,1,1,32]{3,2,1,0} broadcast(multiply.38), dimensions={0,1,2,3}
  reshape.40 = f32[1,32]{1,0} reshape(broadcast.39)
  broadcast.41 = f32[1,64,64,32]{3,2,1,0} broadcast(reshape.40), dimensions={0,3}
  multiply.42 = f32[1,64,64,32]{3,2,1,0} multiply(subtract.34, broadcast.41)
  Arg_2.16 = f32[32]{0} parameter(2)
  reshape.43 = f32[1,1,1,32]{3,2,1,0} reshape(Arg_2.16)
  broadcast.44 = f32[1,1,1,32]{3,2,1,0} broadcast(reshape.43), dimensions={0,1,2,3}
  reshape.45 = f32[1,32]{1,0} reshape(broadcast.44)
  broadcast.46 = f32[1,64,64,32]{3,2,1,0} broadcast(reshape.45), dimensions={0,3}
  add.47 = f32[1,64,64,32]{3,2,1,0} add(multiply.42, broadcast.46)
  ROOT call.48 = f32[1,64,64,32]{3,2,1,0} call(add.47), to_apply=relu.8
}

train_step.49 {
  Arg_0.50 = f32[32]{0} parameter(0)
  Arg_1.51 = f32[32]{0} parameter(1)
  Arg_2.52 = f32[32]{0} parameter(2)
  Arg_3.53 = f32[32]{0} parameter(3)
  Arg_4.54 = f32[32]{0} parameter(4)
  Arg_5.55 = f32[3,3,3,32]{3,2,1,0} parameter(5)
  Arg_6.56 = f32[1,64,64,3]{3,2,1,0} parameter(6)
  ROOT call.57 = f32[1,64,64,32]{3,2,1,0} call(Arg_0.50, Arg_1.51, Arg_2.52, Arg_3.53, Arg_4.54, Arg_5.55, Arg_6.56), to_apply=train_step_0.13
}

ENTRY main.60 {
  Arg_0.1 = f32[32]{0} parameter(0)
  Arg_1.2 = f32[32]{0} parameter(1)
  Arg_2.3 = f32[32]{0} parameter(2)
  Arg_3.4 = f32[32]{0} parameter(3)
  Arg_4.5 = f32[32]{0} parameter(4)
  Arg_5.6 = f32[3,3,3,32]{3,2,1,0} parameter(5)
  Arg_6.7 = f32[1,64,64,3]{3,2,1,0} parameter(6)
  call.58 = f32[1,64,64,32]{3,2,1,0} call(Arg_0.1, Arg_1.2, Arg_2.3, Arg_3.4, Arg_4.5, Arg_5.6, Arg_6.7), to_apply=train_step.49
  ROOT tuple.59 = (f32[1,64,64,32]{3,2,1,0}) tuple(call.58)
}

I see that nn.Conv gets lowered to convolution.23 = f32[1,64,64,32]{3,2,1,0} convolution(Arg_6.20, Arg_5.19), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f but nn.BatchNorm is lowered into a sequence of ops.

I get a sense that it is a translation of implementation here.

Is there some way to avoid this? Especially, since hlo has a native batch_norm_training op?