class_docstr: Tests for methods in piecewise module.
class_docstr_tok: ["Tests","for","methods","in","piecewise","module","."]
docstr: Tests PiecewiseConstantFunc with batching.
docstr_tok: ["Tests","PiecewiseConstantFunc","with","batching","."]
code: def test_piecewise_constant_integral_with_batch(self, dtype):
      """Tests PiecewiseConstantFunc with batching."""
      x = np.array(
          [
              [[0.0, 0.1, 2.0, 11.0], [0.0, 2.0, 3.0, 9.0]],
              [[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]],
          ]
      )
      jump_locations = np.array(
          [[[0.1, 10.0], [1.5, 10.0]], [[1.0, 2.0], [5.0, 6.0]]]
      )
      values = tf.constant(
          [[[3, 4, 5], [3, 4, 5]], [[3, 4, 5], [3, 4, 5]]], dtype=dtype
      )
      piecewise_func = piecewise.PiecewiseConstantFunc(
          jump_locations, values, dtype=dtype
      )
      value = piecewise_func.integrate(x, x + 1.1)
      with self.subTest("Dtype"):
          self.assertEqual(value.dtype.as_numpy_dtype, dtype)
      expected_value = np.array(
          [
              [[4.3, 4.4, 4.4, 5.5], [3.3, 4.4, 4.4, 4.5]],
              [[3.4, 4.5, 5.5, 5.5], [3.4, 4.5, 5.5, 5.5]],
          ]
      )
      self.assertAllClose(value, expected_value, atol=1e-5, rtol=1e-5)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 9 for calling the Piecewise obj's underlying member method code with example usage: obj.test_piecewise_constant_integral_with_batch(dtype) without return types