code: def tensor_array_from_element(elem, size=None, **kwargs):
      """Construct a tf.TensorArray of elements with the dtype + shape of `elem`."""
      if JAX_MODE and jax.dtypes.issubdtype(
          elem.dtype, jax.dtypes.prng_key
      ):
          # If `trace_elt` is a typed prng key, then then it is not possible to create
          # a matching (i.e., with the same custom PRNG) instance/array inside
          # `TensorArray.__init__` given just a `dtype`, `size`, and `shape`.
          #
          # So we cheat -- creating an array of similar keys here.
          kwargs["data"] = _jax_batched_split(elem, size)
      return tf.TensorArray(
          elem.dtype, size=size, element_shape=elem.shape, **kwargs
      )
  
file_path: probability/tensorflow_probability/python/internal/loop_util.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 3 for calling a global function with example usage: tensor_array_from_element(elem, size) and returns: tf