code: def _multiply(tensor, state_sized_tensor, dtype):
      """Multiply `tensor` by a "state sized" tensor and preserve shape."""
      # User should be using a step size that does not alter the state size. This
      # will fail noisily if that is not the case.
      result = tf.cast(tensor, dtype) * tf.cast(
          state_sized_tensor, dtype
      )
      tensorshape_util.set_shape(result, state_sized_tensor.shape)
      return result
  
file_path: probability/tensorflow_probability/python/mcmc/internal/leapfrog_integrator.py
imports: import abc
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 3 for calling a global function with example usage: _multiply(tensor, state_sized_tensor, dtype) and returns: result