code: def __init__(
      self,
      *,
      din: tp.Optional[int] = None,
      mu: float = 0.95,
      rngs: tp.Optional[nnx.Rngs] = None,
  ):
      ...
  
imports: import typing
  import jax
  import jax.numpy
  from jax import random
  from flax.experimental import nnx
label_desc: Use image node_id 3 to create a new BatchNorm object from inherited base classes: nnx with example: obj = BatchNorm()