class_docstr: Tests for methods in piecewise module.
class_docstr_tok: ["Tests","for","methods","in","piecewise","module","."]
docstr: Tests that `values` event shape is `jump_locations` event shape + 1.
docstr_tok: ["Tests","that","`","values","`","event","shape","is","`","jump_locations","`","event","shape","+","1","."]
code: def test_3d_event_shape_with_batch_shape(self, dtype):
      """Tests that `values` event shape is `jump_locations` event shape + 1."""
      x = np.array([[0, 1, 2, 3], [0.5, 1.5, 2.5, 3.5]])
      jump_locations = [[0.5, 2], [0.5, 1.5]]
      values = [
          [[0, 1, 1.5], [2, 3, 0], [1, 0, 1]],
          [[0, 0.5, 1], [1, 3, 2], [2, 3, 1]],
      ]
      piecewise_func = piecewise.PiecewiseConstantFunc(
          jump_locations, values, dtype=dtype
      )
      value = piecewise_func(x)
      integral = piecewise_func.integrate(x, x + 1)
      expected_value = [
          [[0, 1, 1.5], [2, 3, 0], [2, 3, 0], [1, 0, 1]],
          [[0, 0.5, 1], [1, 3, 2], [2, 3, 1], [2, 3, 1]],
      ]
      expected_integral = [
          [[1, 2, 0.75], [2, 3, 0], [1, 0, 1], [1, 0, 1]],
          [[1, 3, 2], [2, 3, 1], [2, 3, 1], [2, 3, 1]],
      ]
      with self.subTest("Values"):
          self.assertAllClose(
              value, expected_value, atol=1e-5, rtol=1e-5
          )
      with self.subTest("Integrals"):
          self.assertAllClose(
              integral, expected_integral, 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 13 for calling the Piecewise obj's underlying member method code with example usage: obj.test_3d_event_shape_with_batch_shape(dtype) without return types