class_docstr: Test for tfp.math.trapz.
class_docstr_tok: ["Test","for","tfp.math.trapz","."]
code: def test_jit(self):
      self.skip_if_no_xla()
  
      def f(y, x=None, dx=None, axis=-1, name=None):
          @tf.function(jit_compile=True)
          def g(y, x=None, dx=None):
              return integration.trapz(
                  y, x=x, dx=dx, axis=axis, name=name
              )
  
          return g(y, x=x, dx=dx)
  
      y = tf.ones((3, 4, 5), dtype=tf.float32, name="x")
      x = tf.ones((3, 4, 5), dtype=tf.float32, name="y")
      dx = tf.constant(0.2, dtype=tf.float32, name="dx")
      axis = tf.constant(1, dtype=tf.int32)
      with self.subTest("f(y)"):
          self.evaluate(f(y))
      with self.subTest("f(y, x)"):
          self.evaluate(f(y, x=x, axis=axis))
      with self.subTest("f(y, dx)"):
          self.evaluate(f(y, dx=dx, axis=axis))
  
imports: import hypothesis
  from hypothesis import strategies
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import hypothesis_testlib
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import integration
label_desc: Use image node_id 15 for calling the TrapzTest obj's underlying member method code with example usage: obj.test_jit() and returns: g, integration