code: def _cond_jax(
      pred, true_fn=None, false_fn=None, name=None
  ):  # pylint: disable=missing-docstring
      from jax import lax  # pylint: disable=g-import-not-at-top
  
      del name
  
      def overridden_true_fn(x):
          del x
          return true_fn()
  
      def overridden_false_fn(x):
          del x
          return false_fn()
  
      return lax.cond(
          pred, None, overridden_true_fn, None, overridden_false_fn
      )
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/control_flow.py
imports: import collections
  import functools
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import dtype
  from tensorflow_probability.python.internal.backend.numpy import nest
  from tensorflow_probability.python.internal.backend.numpy import ops
label_desc: Use image node_id 1 for calling a global function with example usage: _cond_jax(pred, true_fn, false_fn, name) and returns: lax, true_fn, false_fn