code: def _uniform_jax(
      shape,
      minval=0,
      maxval=None,
      dtype=np.float32,
      seed=None,
      name=None,
  ):  # pylint: disable=unused-argument
      """Jax uniform random sampler."""
      import jax.random as jaxrand  # pylint: disable=g-import-not-at-top
  
      if seed is None:
          raise ValueError("Must provide PRNGKey to sample in JAX.")
      dtype = utils.common_dtype([minval, maxval], dtype_hint=dtype)
      final_rank = max(
          [len(shape), len(np.shape(minval)), len(np.shape(maxval))]
      )
      if np.issubdtype(dtype, np.integer):
          if maxval is None:
              raise ValueError(
                  "Must specify maxval for integer dtype {}.".format(
                      dtype
                  )
              )
          shape = _bcast_shape(shape, [minval, maxval])
          # We must match ranks, as lax.max refuses to broadcast different-rank args.
          minval = minval + np.zeros([1] * final_rank, dtype=dtype)
          return jaxrand.randint(
              key=seed,
              shape=shape,
              minval=minval,
              maxval=maxval,
              dtype=dtype,
          )
      else:
          maxval = (
              ops.convert_to_tensor(1, dtype)
              if maxval is None
              else maxval
          )
          shape = _bcast_shape(shape, [minval, maxval])
          # We must match ranks, as lax.max refuses to broadcast different-rank args.
          minval = minval + np.zeros([1] * final_rank, dtype=dtype)
          maxval = maxval + np.zeros([1] * final_rank, dtype=dtype)
          return jaxrand.uniform(
              key=seed,
              shape=shape,
              dtype=dtype,
              minval=minval,
              maxval=maxval,
          )
  
imports: import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
label_desc: Use image node_id 17 for calling a global function with example usage: _uniform_jax(shape, minval, maxval, dtype, seed, name) and returns: jaxrand, jaxrand