code: def _scatter_nd_batch(indices, updates, shape, batch_dims=0):
      """A partial implementation of `scatter_nd` supporting `batch_dims`."""
  
      # `tf.scatter_nd` does not support a `batch_dims` argument.
      # Instead we use the gradient of `tf.gather_nd`.
      # From a purely mathematical perspective this works because
      # (if `tf.scatter_nd` supported `batch_dims`)
      # `gather_nd` and `scatter_nd` (with matching `indices`) are
      # adjoint linear operators and
      # the gradient w.r.t `x` of `dot(y, A(x))` is `adjoint(A)(y)`.
      #
      # Another perspective: back propagating through a "neural" network
      # containing a gather operation carries derivatives backwards through the
      # network, accumulating the derivatives in the locations that
      # were gathered from, ie. they are scattered.
      # If the network multiplies each gathered element by
      # some quantity, then the backwardly propagating derivatives are scaled
      # by this quantity before being scattered.
      # Combining this with the fact that`GradientTape.gradient`
      # starts back-propagation with derivatives equal to `1`, this allows us
      # to use the multipliers to determine the quantities scattered.
      #
      # However, derivatives are only supported for floating point types
      # so we 'tunnel' our types through the `float64` type.
      # So the implmentation is "partial" in the sense that it supports
      # data that can be losslessly converted to `tf.float64` and back.
      dtype = updates.dtype
      internal_dtype = tf.float64
      multipliers = ps.cast(updates, internal_dtype)
  
      def weighted_gathered(zeros):
          return multipliers * tf.gather_nd(
              zeros, indices, batch_dims=batch_dims
          )
  
      zeros = tf.zeros(shape, dtype=internal_dtype)
      _, grad = value_and_gradient(weighted_gathered, zeros)
      return ps.cast(grad, dtype=dtype)
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/weighted_resampling.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.gradient import value_and_gradient
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 6 for calling a global function with example usage: _scatter_nd_batch(indices, updates, shape, batch_dims) and returns: ps, unknown