code: def reset_momentum(momentum, mask):
      new_momentum = (
          momentum if mask is None else momentum * (1.0 - mask)
      )
      return new_momentum
  
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 2 for calling a global function with example usage: reset_momentum(momentum, mask) and returns: new_momentum