code: def test_jax_jit(self):
      if not JAX_MODE:
          self.skipTest("JAX only")
  
      import jax  # pylint: disable=g-import-not-at-top
  
      @jax.jit
      def f(ta, index, value):
          return ta.write(index, value)
  
      ta = nptf.compat.v2.TensorArray(
          dtype=np.float32, size=4, element_shape=(3,)
      )
      ta = f(ta, 1, [2, 3, 4.0])
      ta = f(ta, 3, [4, 5, 6.0])
      self.assertAllEqual(
          [[0, 0, 0], [2, 3, 4], [0, 0, 0], [4, 5, 6.0]], ta.stack()
      )
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/tensor_array_ops_test.py
imports: import numpy
  from tensorflow_probability.python.internal.backend import numpy
  from tensorflow_probability.substrates.numpy.internal import test_util
label_desc: Use image node_id 1 for calling the TensorArrayTest obj's underlying member method code with example usage: obj.test_jax_jit() and returns: ta