docstr: Tests that the values are updated correctly.
docstr_tok: ["Tests","that","the","values","are","updated","correctly","."]
code: def test_maybe_update_along_axis(self, dtype):
      """Tests that the values are updated correctly."""
      tensor = tf.ones([5, 4, 3, 2], dtype=dtype)
      new_tensor = tf.zeros([5, 4, 1, 2], dtype=dtype)
  
      @tf.function
      def maybe_update_along_axis(do_update):
          return utils.maybe_update_along_axis(
              tensor=tensor,
              new_tensor=new_tensor,
              axis=1,
              ind=2,
              do_update=do_update,
          )
  
      updated_tensor = maybe_update_along_axis(True)
      with self.subTest(name="Shape"):
          self.assertEqual(updated_tensor.shape, tensor.shape)
      with self.subTest(name="UpdatedVals"):
          self.assertAllEqual(
              updated_tensor[:, 2, :, :],
              tf.zeros_like(updated_tensor[:, 2, :, :]),
          )
      with self.subTest(name="NotUpdatedVals"):
          self.assertAllEqual(
              updated_tensor[:, 1, :, :],
              tf.ones_like(updated_tensor[:, 2, :, :]),
          )
      with self.subTest(name="DoNotUpdateVals"):
          not_updated_tensor = maybe_update_along_axis(False)
          self.assertAllEqual(not_updated_tensor, tensor)
  
file_path: tf-quant-finance/tf_quant_finance/models/utils_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
  from tf_quant_finance.models import utils
label_desc: Use image node_id 2 for calling the UtilsTest obj's underlying member method code with example usage: obj.test_maybe_update_along_axis(dtype) and returns: utils