docstr: Tests that gradient calculation correctly handles batches.
docstr_tok: ["Tests","that","gradient","calculation","correctly","handles","batches","."]
code: def test_log_soomean_exp_gradient_using_finite_difference_1(self):
      """Tests that gradient calculation correctly handles batches."""
  
      logx_ = np.linspace(-100.0, 100.0, 100).reshape([10, 2, 5])
      logx = tf.constant(logx_)
  
      _, grad_log_soomean_u = self.evaluate(
          gradient.value_and_gradient(
              lambda logx: leave_one_out.log_soomean_exp(logx, axis=0)[
                  0
              ],
              logx,
          )
      )
  
      _, grad_log_mean_u = self.evaluate(
          gradient.value_and_gradient(
              lambda logx: leave_one_out.log_soomean_exp(logx, axis=0)[
                  1
              ],
              logx,
          )
      )
  
      # We skip checking against finite-difference approximation since it
      # doesn't support batches.
  
      # Verify claim in docstring.
      self.assertAllClose(
          np.ones_like(grad_log_mean_u.sum(axis=0)),
          grad_log_mean_u.sum(axis=0),
      )
      self.assertAllClose(
          np.ones_like(grad_log_soomean_u.mean(axis=0)),
          grad_log_soomean_u.mean(axis=0),
      )
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.stats import leave_one_out
label_desc: Use image node_id 8 for calling the LogSooMeanTest obj's underlying member method code with example usage: obj.test_log_soomean_exp_gradient_using_finite_difference_1() without return types