BatchNormΒΆ Auto-generated from flax_nnx_code_defs PyTorchApple MLXFlax NNX PyTorchAPI: torch.nn.BatchNorm2dStrategy: Direct MappingOfficial Docs βApple MLXAPI: mlx.nn.BatchNormStrategy: Direct MappingOfficial Docs βFlax NNXAPI: flax.nnx.BatchNormStrategy: Plugin (batch_norm_unwrap)