code: def gather_precondition(
      precondition_fn: Callable[[GatherArgs], None]
  ):
      """Decorator for specifying a precondition function.
  
      This decorator should be put on a function with argument `arg` of type
      `GatherArgs`. It will first call `precondition_fn` with `arg` (which may throw
      an exception), and then call the function it is decorating with `arg` as well.
      """
  
      def decorator(gather_fn: Callable[[GatherArgs], Any]):
          @wraps(gather_fn)
          def wrapper(args: GatherArgs):
              # Call `precondition_fn`; we assume it may throw an exception.
              precondition_fn(args)
              return gather_fn(args)
  
          return wrapper
  
      return decorator
  
imports: from __future__ import annotations
  import builtins
  from collections.abc import Sequence
  import dataclasses
  from functools import partial, wraps
  import math
  import string
  from typing import Any, Callable, Optional
  from jax._src import core
  from jax import lax
  from jax._src.lax import slicing
  from jax._src import dtypes
  from jax._src import util
  from jax.experimental.jax2tf import jax2tf
  import numpy
  import tensorflow
label_desc: Use image node_id 25 for calling a global function with example usage: gather_precondition(precondition_fn) and returns: decorator, wrapper, gather_fn