code: def weight_reinit_zero(param, mask):
      if mask is None:
          return param
      else:
          new_param = jnp.zeros_like(param)
          param = jnp.where(mask == 1, new_param, param)
          return param
  
imports: import functools
  import logging
  import flax
  from flax import linen
  import gin
  import jax
  from jax import random
  import jax.numpy
  import optax
label_desc: Use image node_id 3 for calling a global function with example usage: weight_reinit_zero(param, mask) and returns: param, param