class_docstr: Runs one step of Uncalibrated Langevin discretized diffusion.
  
  The class generates a Langevin proposal using `_euler_method` function and
  also computes helper `UncalibratedLangevinKernelResults` for the next
  iteration.
  
  Warning: this kernel will not result in a chain which converges to the
  `target_log_prob`. To get a convergent MCMC, use
  `MetropolisAdjustedLangevinAlgorithm(...)` or
  `MetropolisHastings(UncalibratedLangevin(...))`.
  
  For more details on `UncalibratedLangevin`, see
  `MetropolisAdjustedLangevinAlgorithm`.
class_docstr_tok: ["Runs","one","step","of","Uncalibrated","Langevin","discretized","diffusion",".","The","class","generates","a","Langevin","proposal","using","`","_euler_method","`","function","and","also","computes","helper","`","UncalibratedLangevinKernelResults","`","for","the","next","iteration",".","Warning",":","this","kernel","will","not","result","in","a","chain","which","converges","to","the","`","target_log_prob","`",".","To","get","a","convergent","MCMC",",","use","`","MetropolisAdjustedLangevinAlgorithm","(","...",")","`","or","`","MetropolisHastings","(","UncalibratedLangevin","(","...",")",")","`",".","For","more","details","on","`","UncalibratedLangevin","`",",","see","`","MetropolisAdjustedLangevinAlgorithm","`","."]
code: def one_step(self, current_state, previous_kernel_results, seed=None):
      with tf.name_scope(
          mcmc_util.make_name(self.name, "mala", "one_step")
      ):
          with tf.name_scope("initialize"):
              # Prepare input arguments to be passed to `_euler_method`.
              [
                  current_state_parts,
                  step_size_parts,
                  current_target_log_prob,
                  _,  # grads_target_log_prob
                  current_volatility_parts,
                  _,  # grads_volatility
                  current_drift_parts,
              ] = _prepare_args(
                  self.target_log_prob_fn,
                  self.volatility_fn,
                  current_state,
                  self.step_size,
                  previous_kernel_results.target_log_prob,
                  previous_kernel_results.grads_target_log_prob,
                  previous_kernel_results.volatility,
                  previous_kernel_results.grads_volatility,
                  previous_kernel_results.diffusion_drift,
                  self.parallel_iterations,
              )
  
              seed = samplers.sanitize_seed(
                  seed
              )  # Retain for diagnostics.
              seeds = list(
                  samplers.split_seed(
                      seed,
                      n=len(current_state_parts),
                      salt="langevin.one_step",
                  )
              )
              seeds = distribute_lib.fold_in_axis_index(
                  seeds, self.experimental_shard_axis_names
              )
  
              random_draw_parts = []
              for state_part, part_seed in zip(
                  current_state_parts, seeds
              ):
                  random_draw_parts.append(
                      samplers.normal(
                          shape=ps.shape(state_part),
                          dtype=dtype_util.base_dtype(state_part.dtype),
                          seed=part_seed,
                      )
                  )
  
          # Number of independent chains run by the algorithm.
          independent_chain_ndims = ps.rank(current_target_log_prob)
  
          # Generate the next state of the algorithm using Euler-Maruyama method.
          next_state_parts = _euler_method(
              random_draw_parts,
              current_state_parts,
              current_drift_parts,
              step_size_parts,
              current_volatility_parts,
          )
  
          # Compute helper `UncalibratedLangevinKernelResults` to be processed by
          # `_compute_log_acceptance_correction` and in the next iteration of
          # `one_step` function.
          [
              _,  # state_parts
              _,  # step_sizes
              next_target_log_prob,
              next_grads_target_log_prob,
              next_volatility_parts,
              next_grads_volatility,
              next_drift_parts,
          ] = _prepare_args(
              self.target_log_prob_fn,
              self.volatility_fn,
              next_state_parts,
              step_size_parts,
              parallel_iterations=self.parallel_iterations,
          )
  
          def maybe_flatten(x):
              return (
                  x if mcmc_util.is_list_like(current_state) else x[0]
              )
  
          # Decide whether to compute the acceptance ratio
          log_acceptance_correction_compute = _compute_log_acceptance_correction(
              current_state_parts,
              next_state_parts,
              current_volatility_parts,
              next_volatility_parts,
              current_drift_parts,
              next_drift_parts,
              step_size_parts,
              independent_chain_ndims,
              experimental_shard_axis_names=self.experimental_shard_axis_names,
          )
          log_acceptance_correction_skip = tf.zeros_like(
              next_target_log_prob
          )
  
          log_acceptance_correction = tf.cond(
              pred=self.compute_acceptance,
              true_fn=lambda: log_acceptance_correction_compute,
              false_fn=lambda: log_acceptance_correction_skip,
          )
  
          return [
              maybe_flatten(next_state_parts),
              UncalibratedLangevinKernelResults(
                  log_acceptance_correction=log_acceptance_correction,
                  target_log_prob=next_target_log_prob,
                  grads_target_log_prob=next_grads_target_log_prob,
                  volatility=maybe_flatten(next_volatility_parts),
                  grads_volatility=next_grads_volatility,
                  diffusion_drift=next_drift_parts,
                  seed=seed,
              ),
          ]
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 10 for calling the UncalibratedLangevin obj's underlying member method code with example usage: obj.one_step(current_state, previous_kernel_results, seed) and returns: list, unknown