code: def _get_zero_fn(
      mask: jnp.ndarray,
  ) -> Callable[[jnp.ndarray], jnp.ndarray]:
      return lambda x: _expand_trailing_dimensions(mask, x) * x
  
imports: import functools
  from typing import Any, Callable, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Union
  import jax
  from jax import lax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  import numpy
label_desc: Use image node_id 35 for calling a global function with example usage: _get_zero_fn(mask) and returns: function