code: # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """LSTM Block Cell ops."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  
  from tensorflow.contrib.rnn.ops import gen_lstm_ops
  from tensorflow.contrib.util import loader
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.layers import base as base_layer
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import rnn_cell_impl
  from tensorflow.python.platform import resource_loader
  
  _lstm_ops_so = loader.load_op_library(
      resource_loader.get_path_to_datafile("_lstm_ops.so"))
  
  LayerRNNCell = rnn_cell_impl.LayerRNNCell  # pylint: disable=invalid-name
  
  
  # pylint: disable=invalid-name
  def _lstm_block_cell(x,
                       cs_prev,
                       h_prev,
                       w,
                       b,
                       wci=None,
                       wcf=None,
                       wco=None,
                       forget_bias=None,
                       cell_clip=None,
                       use_peephole=None,
                       name=None):
    r"""Computes the LSTM cell forward propagation for 1 time step.
  
    This implementation uses 1 weight matrix and 1 bias vector, and there's an
    optional peephole connection.
  
    This kernel op implements the following mathematical equations:
  
    ```python
    xh = [x, h_prev]
    [i, ci, f, o] = xh * w + b
    f = f + forget_bias
  
    if not use_peephole:
      wci = wcf = wco = 0
  
    i = sigmoid(cs_prev * wci + i)
    f = sigmoid(cs_prev * wcf + f)
    ci = tanh(ci)
  
    cs = ci .* i + cs_prev .* f
    cs = clip(cs, cell_clip)
  
    o = sigmoid(cs * wco + o)
    co = tanh(cs)
    h = co .* o
    ```
  
    Args:
      x: A `Tensor`. Must be one of the following types: `float32`.
        The input to the LSTM cell, shape (batch_size, num_inputs).
      cs_prev: A `Tensor`. Must have the same type as `x`.
        Value of the cell state at previous time step.
      h_prev: A `Tensor`. Must have the same type as `x`.
        Output of the previous cell at previous time step.
      w: A `Tensor`. Must have the same type as `x`. The weight matrix.
      b: A `Tensor`. Must have the same type as `x`. The bias vector.
      wci: A `Tensor`. Must have the same type as `x`.
        The weight matrix for input gate peephole connection.
      wcf: A `Tensor`. Must have the same type as `x`.
        The weight matrix for forget gate peephole connection.
      wco: A `Tensor`. Must have the same type as `x`.
        The weight matrix for output gate peephole connection.
      forget_bias: An optional `float`. Defaults to `1`. The forget gate bias.
      cell_clip: An optional `float`. Defaults to `-1` (no clipping).
        Value to clip the 'cs' value to. Disable by setting to negative value.
      use_peephole: An optional `bool`. Defaults to `False`.
        Whether to use peephole weights.
      name: A name for the operation (optional).
  
    Returns:
      A tuple of `Tensor` objects (i, cs, f, o, ci, co, h).
      i: A `Tensor`. Has the same type as `x`. The input gate.
      cs: A `Tensor`. Has the same type as `x`. The cell state before the tanh.
      f: A `Tensor`. Has the same type as `x`. The forget gate.
      o: A `Tensor`. Has the same type as `x`. The output gate.
      ci: A `Tensor`. Has the same type as `x`. The cell input.
      co: A `Tensor`. Has the same type as `x`. The cell after the tanh.
      h: A `Tensor`. Has the same type as `x`. The output h vector.
  
    Raises:
      ValueError: If cell_size is None.
    """
    if wci is None:
      cell_size = cs_prev.get_shape().with_rank(2)[1].value
      if cell_size is None:
        raise ValueError("cell_size from `cs_prev` should not be None.")
      wci = array_ops.constant(0, dtype=dtypes.float32, shape=[cell_size])
      wcf = wci
      wco = wci
  
    # pylint: disable=protected-access
    return gen_lstm_ops.lstm_block_cell(
        x=x,
        cs_prev=cs_prev,
        h_prev=h_prev,
        w=w,
        wci=wci,
        wcf=wcf,
        wco=wco,
        b=b,
        forget_bias=forget_bias,
        cell_clip=cell_clip if cell_clip is not None else -1,
        use_peephole=use_peephole,
        name=name)
    # pylint: enable=protected-access
  
  
  def _block_lstm(seq_len_max,
                  x,
                  w,
                  b,
                  cs_prev=None,
                  h_prev=None,
                  wci=None,
                  wcf=None,
                  wco=None,
                  forget_bias=None,
                  cell_clip=None,
                  use_peephole=None,
                  name=None):
    r"""TODO(williamchan): add doc.
  
    Args:
      seq_len_max: A `Tensor` of type `int64`.
      x: A list of at least 1 `Tensor` objects of the same type in: `float32`.
      w: A `Tensor`. Must have the same type as `x`.
      b: A `Tensor`. Must have the same type as `x`.
      cs_prev: A `Tensor`. Must have the same type as `x`.
      h_prev: A `Tensor`. Must have the same type as `x`.
      wci: A `Tensor`. Must have the same type as `x`.
      wcf: A `Tensor`. Must have the same type as `x`.
      wco: A `Tensor`. Must have the same type as `x`.
      forget_bias: An optional `float`. Defaults to `1`.
      cell_clip: An optional `float`. Defaults to `-1` (no clipping).
      use_peephole: An optional `bool`. Defaults to `False`.
      name: A name for the operation (optional).
  
    Returns:
      A tuple of `Tensor` objects (i, cs, f, o, ci, co, h).
      i: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
      cs: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
      f: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
      o: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
      ci: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
      co: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
      h: A list with the same number of `Tensor` objects as `x` of `Tensor`
      objects of the same type as x.
  
    Raises:
      ValueError: If `b` does not have a valid shape.
    """
    batch_size = x[0].get_shape().with_rank(2)[0].value
    cell_size4 = b.get_shape().with_rank(1)[0].value
    if cell_size4 is None:
      raise ValueError("`b` shape must not be None.")
    cell_size = cell_size4 / 4
    zero_state = None
    if cs_prev is None or h_prev is None:
      zero_state = array_ops.constant(
          0, dtype=dtypes.float32, shape=[batch_size, cell_size])
    if cs_prev is None:
      cs_prev = zero_state
    if h_prev is None:
      h_prev = zero_state
    if wci is None:
      wci = array_ops.constant(0, dtype=dtypes.float32, shape=[cell_size])
      wcf = wci
      wco = wci
  
    # pylint: disable=protected-access
    i, cs, f, o, ci, co, h = gen_lstm_ops.block_lstm(
        seq_len_max=seq_len_max,
        x=array_ops.stack(x),
        cs_prev=cs_prev,
        h_prev=h_prev,
        w=w,
        wci=wci,
        wcf=wcf,
        wco=wco,
        b=b,
        forget_bias=forget_bias,
        cell_clip=cell_clip if cell_clip is not None else -1,
        name=name,
        use_peephole=use_peephole)
  
    return array_ops.unstack(i), array_ops.unstack(cs), array_ops.unstack(
        f), array_ops.unstack(o), array_ops.unstack(ci), array_ops.unstack(
            co), array_ops.unstack(h)
    # pylint: enable=protected-access
    # pylint: enable=invalid-name
  
  
  _lstm_block_cell_grad_outputs = ["cs_prev_grad", "dicfo"]
  
  
  @ops.RegisterGradient("LSTMBlockCell")
  def _LSTMBlockCellGrad(op, *grad):
    """Gradient for LSTMBlockCell."""
    (x, cs_prev, h_prev, w, wci, wcf, wco, b) = op.inputs
    (i, cs, f, o, ci, co, _) = op.outputs
    (_, cs_grad, _, _, _, _, h_grad) = grad
  
    batch_size = x.get_shape().with_rank(2)[0].value
    if batch_size is None:
      batch_size = -1
    input_size = x.get_shape().with_rank(2)[1].value
    if input_size is None:
      raise ValueError("input_size from `x` should not be None.")
    cell_size = cs_prev.get_shape().with_rank(2)[1].value
    if cell_size is None:
      raise ValueError("cell_size from `cs_prev` should not be None.")
  
    (cs_prev_grad, dicfo, wci_grad, wcf_grad,
     wco_grad) = gen_lstm_ops.lstm_block_cell_grad(
         x,
         cs_prev,
         h_prev,
         w,
         wci,
         wcf,
         wco,
         b,
         i,
         cs,
         f,
         o,
         ci,
         co,
         cs_grad,
         h_grad,
         use_peephole=op.get_attr("use_peephole"))
  
    # Backprop from dicfo to xh.
    xh_grad = math_ops.matmul(dicfo, w, transpose_b=True)
  
    x_grad = array_ops.slice(xh_grad, (0, 0), (batch_size, input_size))
    x_grad.get_shape().merge_with(x.get_shape())
  
    h_prev_grad = array_ops.slice(xh_grad, (0, input_size),
                                  (batch_size, cell_size))
    h_prev_grad.get_shape().merge_with(h_prev.get_shape())
  
    # Backprop from dicfo to w.
    xh = array_ops.concat([x, h_prev], 1)
    w_grad = math_ops.matmul(xh, dicfo, transpose_a=True)
    w_grad.get_shape().merge_with(w.get_shape())
  
    # Backprop from dicfo to b.
    b_grad = nn_ops.bias_add_grad(dicfo)
    b_grad.get_shape().merge_with(b.get_shape())
  
    return (x_grad, cs_prev_grad, h_prev_grad, w_grad, wci_grad, wcf_grad,
            wco_grad, b_grad)
  
  
  @ops.RegisterGradient("BlockLSTM")
  def _BlockLSTMGrad(op, *grad):
    """Gradient for BlockLSTM."""
    seq_len_max, x, cs_prev, h_prev, w, wci, wcf, wco, b = op.inputs
    i, cs, f, o, ci, co, h = op.outputs
  
    cs_grad = grad[1]
    h_grad = grad[6]
  
    (x_grad, cs_prev_grad, h_prev_grad, w_grad, wci_grad, wcf_grad, wco_grad,
     b_grad) = gen_lstm_ops.block_lstm_grad(
         seq_len_max,
         x,
         cs_prev,
         h_prev,
         w,
         wci,
         wcf,
         wco,
         b,
         i,
         cs,
         f,
         o,
         ci,
         co,
         h,
         cs_grad,
         h_grad,
         use_peephole=op.get_attr("use_peephole"))
  
    return [
        None, x_grad, cs_prev_grad, h_prev_grad, w_grad, wci_grad, wcf_grad,
        wco_grad, b_grad
    ]
  
  
  class LSTMBlockCell(LayerRNNCell):
    """Basic LSTM recurrent network cell.
  
    The implementation is based on: http://arxiv.org/abs/1409.2329.
  
    We add `forget_bias` (default: 1) to the biases of the forget gate in order to
    reduce the scale of forgetting in the beginning of the training.
  
    Unlike `rnn_cell_impl.LSTMCell`, this is a monolithic op and should be much
    faster.  The weight and bias matrices should be compatible as long as the
    variable scope matches.
    """
  
    def __init__(self,
                 num_units,
                 forget_bias=1.0,
                 cell_clip=None,
                 use_peephole=False,
                 reuse=None,
                 name="lstm_cell"):
      """Initialize the basic LSTM cell.
  
      Args:
        num_units: int, The number of units in the LSTM cell.
        forget_bias: float, The bias added to forget gates (see above).
        cell_clip: An optional `float`. Defaults to `-1` (no clipping).
        use_peephole: Whether to use peephole connections or not.
        reuse: (optional) boolean describing whether to reuse variables in an
          existing scope.  If not `True`, and the existing scope already has the
          given variables, an error is raised.
        name: String, the name of the layer. Layers with the same name will
          share weights, but to avoid mistakes we require reuse=True in such
          cases.  By default this is "lstm_cell", for variable-name compatibility
          with `tf.nn.rnn_cell.LSTMCell`.
  
        When restoring from CudnnLSTM-trained checkpoints, must use
        CudnnCompatibleLSTMBlockCell instead.
      """
      super(LSTMBlockCell, self).__init__(_reuse=reuse, name=name)
      self._num_units = num_units
      self._forget_bias = forget_bias
      self._use_peephole = use_peephole
      self._cell_clip = cell_clip if cell_clip is not None else -1
      self._names = {
          "W": "kernel",
          "b": "bias",
          "wci": "w_i_diag",
          "wcf": "w_f_diag",
          "wco": "w_o_diag",
          "scope": "lstm_cell"
      }
      # Inputs must be 2-dimensional.
      self.input_spec = base_layer.InputSpec(ndim=2)
  
    @property
    def state_size(self):
      return rnn_cell_impl.LSTMStateTuple(self._num_units, self._num_units)
  
    @property
    def output_size(self):
      return self._num_units
  
    def build(self, inputs_shape):
      if not inputs_shape[1].value:
        raise ValueError(
            "Expecting inputs_shape[1] to be set: %s" % str(inputs_shape))
      input_size = inputs_shape[1].value
      self._kernel = self.add_variable(
          self._names["W"], [input_size + self._num_units, self._num_units * 4])
      self._bias = self.add_variable(
          self._names["b"], [self._num_units * 4],
          initializer=init_ops.constant_initializer(0.0))
      if self._use_peephole:
        self._w_i_diag = self.add_variable(self._names["wci"], [self._num_units])
        self._w_f_diag = self.add_variable(self._names["wcf"], [self._num_units])
        self._w_o_diag = self.add_variable(self._names["wco"], [self._num_units])
  
      self.built = True
  
    def call(self, inputs, state):
      """Long short-term memory cell (LSTM)."""
      if len(state) != 2:
        raise ValueError("Expecting state to be a tuple with length 2.")
  
      if self._use_peephole:
        wci = self._w_i_diag
        wcf = self._w_f_diag
        wco = self._w_o_diag
      else:
        wci = wcf = wco = array_ops.zeros([self._num_units])
  
      (cs_prev, h_prev) = state
      (_, cs, _, _, _, _, h) = _lstm_block_cell(
          inputs,
          cs_prev,
          h_prev,
          self._kernel,
          self._bias,
          wci=wci,
          wcf=wcf,
          wco=wco,
          forget_bias=self._forget_bias,
          cell_clip=self._cell_clip,
          use_peephole=self._use_peephole)
  
      new_state = rnn_cell_impl.LSTMStateTuple(cs, h)
      return h, new_state
  
  
  class LSTMBlockWrapper(base_layer.Layer):
    """This is a helper class that provides housekeeping for LSTM cells.
  
    This may be useful for alternative LSTM and similar type of cells.
    The subclasses must implement `_call_cell` method and `num_units` property.
    """
  
    @abc.abstractproperty
    def num_units(self):
      """Number of units in this cell (output dimension)."""
      pass
  
    @abc.abstractmethod
    def _call_cell(self, inputs, initial_cell_state, initial_output, dtype,
                   sequence_length):
      """Run this LSTM on inputs, starting from the given state.
  
      This method must be implemented by subclasses and does the actual work
      of calling the cell.
  
      Args:
        inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]`
        initial_cell_state: initial value for cell state, shape `[batch_size,
          self._num_units]`
        initial_output: initial value of cell output, shape `[batch_size,
          self._num_units]`
        dtype: The data type for the initial state and expected output.
        sequence_length: Specifies the length of each sequence in inputs. An int32
          or int64 vector (tensor) size [batch_size], values in [0, time_len) or
            None.
  
      Returns:
        A pair containing:
  
        - State: A `3-D` tensor of shape `[time_len, batch_size, output_size]`
        - Output: A `3-D` tensor of shape `[time_len, batch_size, output_size]`
      """
      pass
  
    def call(self, inputs, initial_state=None, dtype=None, sequence_length=None):
      """Run this LSTM on inputs, starting from the given state.
  
      Args:
        inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]`
          or a list of `time_len` tensors of shape `[batch_size, input_size]`.
        initial_state: a tuple `(initial_cell_state, initial_output)` with tensors
          of shape `[batch_size, self._num_units]`. If this is not provided, the
          cell is expected to create a zero initial state of type `dtype`.
        dtype: The data type for the initial state and expected output. Required
          if `initial_state` is not provided or RNN state has a heterogeneous
          dtype.
        sequence_length: Specifies the length of each sequence in inputs. An
          `int32` or `int64` vector (tensor) size `[batch_size]`, values in `[0,
          time_len).`
          Defaults to `time_len` for each element.
  
      Returns:
        A pair containing:
  
        - Output: A `3-D` tensor of shape `[time_len, batch_size, output_size]`
          or a list of time_len tensors of shape `[batch_size, output_size]`,
          to match the type of the `inputs`.
        - Final state: a tuple `(cell_state, output)` matching `initial_state`.
  
      Raises:
        ValueError: in case of shape mismatches
      """
      is_list = isinstance(inputs, list)
      if is_list:
        inputs = array_ops.stack(inputs)
      inputs_shape = inputs.get_shape().with_rank(3)
      if not inputs_shape[2]:
        raise ValueError("Expecting inputs_shape[2] to be set: %s" % inputs_shape)
      batch_size = inputs_shape[1].value
      if batch_size is None:
        batch_size = array_ops.shape(inputs)[1]
      time_len = inputs_shape[0].value
      if time_len is None:
        time_len = array_ops.shape(inputs)[0]
  
      # Provide default values for initial_state and dtype
      if initial_state is None:
        if dtype is None:
          raise ValueError("Either initial_state or dtype needs to be specified")
        z = array_ops.zeros(
            array_ops.stack([batch_size, self.num_units]), dtype=dtype)
        initial_state = z, z
      else:
        if len(initial_state) != 2:
          raise ValueError(
              "Expecting initial_state to be a tuple with length 2 or None")
        if dtype is None:
          dtype = initial_state[0].dtype
  
      # create the actual cell
      if sequence_length is not None:
        sequence_length = ops.convert_to_tensor(sequence_length)
      initial_cell_state, initial_output = initial_state  # pylint: disable=unpacking-non-sequence
      cell_states, outputs = self._call_cell(
          inputs, initial_cell_state, initial_output, dtype, sequence_length)
  
      if sequence_length is not None:
        # Mask out the part beyond sequence_length
        mask = array_ops.transpose(
            array_ops.sequence_mask(sequence_length, time_len, dtype=dtype),
            [1, 0])
        mask = array_ops.tile(
            array_ops.expand_dims(mask, [-1]), [1, 1, self.num_units])
        outputs *= mask
        # Prepend initial states to cell_states and outputs for indexing to work
        # correctly,since we want to access the last valid state at
        # sequence_length - 1, which can even be -1, corresponding to the
        # initial state.
        mod_cell_states = array_ops.concat(
            [array_ops.expand_dims(initial_cell_state, [0]), cell_states], 0)
        mod_outputs = array_ops.concat(
            [array_ops.expand_dims(initial_output, [0]), outputs], 0)
        final_cell_state = self._gather_states(mod_cell_states, sequence_length,
                                               batch_size)
        final_output = self._gather_states(mod_outputs, sequence_length,
                                           batch_size)
      else:
        # No sequence_lengths used: final state is the last state
        final_cell_state = cell_states[-1]
        final_output = outputs[-1]
  
      if is_list:
        # Input was a list, so return a list
        outputs = array_ops.unstack(outputs)
  
      final_state = rnn_cell_impl.LSTMStateTuple(final_cell_state, final_output)
      return outputs, final_state
  
    def _gather_states(self, data, indices, batch_size):
      """Produce `out`, s.t. out(i, j) = data(indices(i), i, j)."""
      mod_indices = indices * batch_size + math_ops.range(batch_size)
      return array_ops.gather(
          array_ops.reshape(data, [-1, self.num_units]), mod_indices)
  
  
  class LSTMBlockFusedCell(LSTMBlockWrapper):
    """FusedRNNCell implementation of LSTM.
  
    This is an extremely efficient LSTM implementation, that uses a single TF op
    for the entire LSTM. It should be both faster and more memory-efficient than
    LSTMBlockCell defined above.
  
    The implementation is based on: http://arxiv.org/abs/1409.2329.
  
    We add forget_bias (default: 1) to the biases of the forget gate in order to
    reduce the scale of forgetting in the beginning of the training.
  
    The variable naming is consistent with `rnn_cell_impl.LSTMCell`.
    """
  
    def __init__(self,
                 num_units,
                 forget_bias=1.0,
                 cell_clip=None,
                 use_peephole=False,
                 reuse=None,
                 name="lstm_fused_cell"):
      """Initialize the LSTM cell.
  
      Args:
        num_units: int, The number of units in the LSTM cell.
        forget_bias: float, The bias added to forget gates (see above).
        cell_clip: clip the cell to this value. Default is no cell clipping.
        use_peephole: Whether to use peephole connections or not.
        reuse: (optional) boolean describing whether to reuse variables in an
          existing scope.  If not `True`, and the existing scope already has the
          given variables, an error is raised.
        name: String, the name of the layer. Layers with the same name will
          share weights, but to avoid mistakes we require reuse=True in such
          cases.  By default this is "lstm_cell", for variable-name compatibility
          with `tf.nn.rnn_cell.LSTMCell`.
      """
      super(LSTMBlockFusedCell, self).__init__(_reuse=reuse, name=name)
      self._num_units = num_units
      self._forget_bias = forget_bias
      self._cell_clip = cell_clip if cell_clip is not None else -1
      self._use_peephole = use_peephole
  
      # Inputs must be 3-dimensional.
      self.input_spec = base_layer.InputSpec(ndim=3)
  
    @property
    def num_units(self):
      """Number of units in this cell (output dimension)."""
      return self._num_units
  
    def build(self, input_shape):
      input_size = input_shape[2].value
      self._kernel = self.add_variable(
          "kernel", [input_size + self._num_units, self._num_units * 4])
      self._bias = self.add_variable(
          "bias", [self._num_units * 4],
          initializer=init_ops.constant_initializer(0.0))
      if self._use_peephole:
        self._w_i_diag = self.add_variable("w_i_diag", [self._num_units])
        self._w_f_diag = self.add_variable("w_f_diag", [self._num_units])
        self._w_o_diag = self.add_variable("w_o_diag", [self._num_units])
  
      self.built = True
  
    def _call_cell(self,
                   inputs,
                   initial_cell_state=None,
                   initial_output=None,
                   dtype=None,
                   sequence_length=None):
      """Run this LSTM on inputs, starting from the given state.
  
      Args:
        inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]`
        initial_cell_state: initial value for cell state, shape `[batch_size,
          self._num_units]`
        initial_output: initial value of cell output, shape `[batch_size,
          self._num_units]`
        dtype: The data type for the initial state and expected output.
        sequence_length: Specifies the length of each sequence in inputs. An
          `int32` or `int64` vector (tensor) size `[batch_size]`, values in `[0,
          time_len)` or None.
  
      Returns:
        A pair containing:
  
        - Cell state (cs): A `3-D` tensor of shape `[time_len, batch_size,
                           output_size]`
        - Output (h): A `3-D` tensor of shape `[time_len, batch_size,
                      output_size]`
      """
  
      inputs_shape = inputs.get_shape().with_rank(3)
      time_len = inputs_shape[0].value
      if time_len is None:
        time_len = array_ops.shape(inputs)[0]
  
      if self._use_peephole:
        wci = self._w_i_diag
        wco = self._w_o_diag
        wcf = self._w_f_diag
      else:
        wci = wcf = wco = array_ops.zeros([self._num_units], dtype=dtype)
  
      if sequence_length is None:
        max_seq_len = math_ops.to_int64(time_len)
      else:
        max_seq_len = math_ops.to_int64(math_ops.reduce_max(sequence_length))
  
      _, cs, _, _, _, _, h = gen_lstm_ops.block_lstm(
          seq_len_max=max_seq_len,
          x=inputs,
          cs_prev=initial_cell_state,
          h_prev=initial_output,
          w=self._kernel,
          wci=wci,
          wcf=wcf,
          wco=wco,
          b=self._bias,
          forget_bias=self._forget_bias,
          cell_clip=self._cell_clip,
          use_peephole=self._use_peephole)
      return cs, h
   # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Hamiltonian Monte Carlo, a gradient-based MCMC algorithm.
  
  @@sample_chain
  @@sample_annealed_importance_chain
  @@kernel
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  import numpy as np
  
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import functional_ops
  from tensorflow.python.ops import gradients_impl as gradients_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops.distributions import util as distributions_util
  
  __all__ = [
      "sample_chain",
      "sample_annealed_importance_chain",
      "kernel",
  ]
  
  
  KernelResults = collections.namedtuple(
      "KernelResults",
      [
          "acceptance_probs",
          "current_grads_target_log_prob",  # "Current result" means "accepted".
          "current_target_log_prob",  # "Current result" means "accepted".
          "energy_change",
          "is_accepted",
          "proposed_grads_target_log_prob",
          "proposed_state",
          "proposed_target_log_prob",
          "random_positive",
      ])
  
  
  def _make_dummy_kernel_results(
      dummy_state,
      dummy_target_log_prob,
      dummy_grads_target_log_prob):
    return KernelResults(
        acceptance_probs=dummy_target_log_prob,
        current_grads_target_log_prob=dummy_grads_target_log_prob,
        current_target_log_prob=dummy_target_log_prob,
        energy_change=dummy_target_log_prob,
        is_accepted=array_ops.ones_like(dummy_target_log_prob, dtypes.bool),
        proposed_grads_target_log_prob=dummy_grads_target_log_prob,
        proposed_state=dummy_state,
        proposed_target_log_prob=dummy_target_log_prob,
        random_positive=dummy_target_log_prob,
    )
  
  
  def sample_chain(
      num_results,
      target_log_prob_fn,
      current_state,
      step_size,
      num_leapfrog_steps,
      num_burnin_steps=0,
      num_steps_between_results=0,
      seed=None,
      current_target_log_prob=None,
      current_grads_target_log_prob=None,
      name=None):
    """Runs multiple iterations of one or more Hamiltonian Monte Carlo chains.
  
    Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) algorithm
    that takes a series of gradient-informed steps to produce a Metropolis
    proposal. This function samples from an HMC Markov chain at `current_state`
    and whose stationary distribution has log-unnormalized-density
    `target_log_prob_fn()`.
  
    This function samples from multiple chains in parallel. It assumes that the
    the leftmost dimensions of (each) `current_state` (part) index an independent
    chain.  The function `target_log_prob_fn()` sums log-probabilities across
    event dimensions (i.e., current state (part) rightmost dimensions). Each
    element of the output of `target_log_prob_fn()` represents the (possibly
    unnormalized) log-probability of the joint distribution over (all) the current
    state (parts).
  
    The `current_state` can be represented as a single `Tensor` or a `list` of
    `Tensors` which collectively represent the current state. When specifying a
    `list`, one must also specify a list of `step_size`s.
  
    Note: `target_log_prob_fn` is called exactly twice.
  
    Only one out of every `num_steps_between_samples + 1` steps is included in the
    returned results. This "thinning" comes at a cost of reduced statistical
    power, while reducing memory requirements and autocorrelation. For more
    discussion see [1].
  
    [1]: "Statistically efficient thinning of a Markov chain sampler."
         Art B. Owen. April 2017.
         http://statweb.stanford.edu/~owen/reports/bestthinning.pdf
  
    #### Examples:
  
    ##### Sample from a diagonal-variance Gaussian.
  
    ```python
    tfd = tf.contrib.distributions
  
    def make_likelihood(true_variances):
      return tfd.MultivariateNormalDiag(
          scale_diag=tf.sqrt(true_variances))
  
    dims = 10
    dtype = np.float32
    true_variances = tf.linspace(dtype(1), dtype(3), dims)
    likelihood = make_likelihood(true_variances)
  
    states, kernel_results = hmc.sample_chain(
        num_results=1000,
        target_log_prob_fn=likelihood.log_prob,
        current_state=tf.zeros(dims),
        step_size=0.5,
        num_leapfrog_steps=2,
        num_burnin_steps=500)
  
    # Compute sample stats.
    sample_mean = tf.reduce_mean(states, axis=0)
    sample_var = tf.reduce_mean(
        tf.squared_difference(states, sample_mean),
        axis=0)
    ```
  
    ##### Sampling from factor-analysis posteriors with known factors.
  
    I.e.,
  
    ```none
    for i=1..n:
      w[i] ~ Normal(0, eye(d))            # prior
      x[i] ~ Normal(loc=matmul(w[i], F))  # likelihood
    ```
  
    where `F` denotes factors.
  
    ```python
    tfd = tf.contrib.distributions
  
    def make_prior(dims, dtype):
      return tfd.MultivariateNormalDiag(
          loc=tf.zeros(dims, dtype))
  
    def make_likelihood(weights, factors):
      return tfd.MultivariateNormalDiag(
          loc=tf.tensordot(weights, factors, axes=[[0], [-1]]))
  
    # Setup data.
    num_weights = 10
    num_factors = 4
    num_chains = 100
    dtype = np.float32
  
    prior = make_prior(num_weights, dtype)
    weights = prior.sample(num_chains)
    factors = np.random.randn(num_factors, num_weights).astype(dtype)
    x = make_likelihood(weights, factors).sample(num_chains)
  
    def target_log_prob(w):
      # Target joint is: `f(w) = p(w, x | factors)`.
      return prior.log_prob(w) + make_likelihood(w, factors).log_prob(x)
  
    # Get `num_results` samples from `num_chains` independent chains.
    chains_states, kernels_results = hmc.sample_chain(
        num_results=1000,
        target_log_prob_fn=target_log_prob,
        current_state=tf.zeros([num_chains, dims], dtype),
        step_size=0.1,
        num_leapfrog_steps=2,
        num_burnin_steps=500)
  
    # Compute sample stats.
    sample_mean = tf.reduce_mean(chains_states, axis=[0, 1])
    sample_var = tf.reduce_mean(
        tf.squared_difference(chains_states, sample_mean),
        axis=[0, 1])
    ```
  
    Args:
      num_results: Integer number of Markov chain draws.
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step size
        for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but too-large
        step sizes make rejection exponentially more likely. When possible, it's
        often helpful to match per-variable step sizes to the standard deviations
        of the target distribution in each variable.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to `step_size *
        num_leapfrog_steps`.
      num_burnin_steps: Integer number of chain steps to take before starting to
        collect results.
        Default value: 0 (i.e., no burn-in).
      num_steps_between_results: Integer number of chain steps between collecting
        a result. Only one out of every `num_steps_between_samples + 1` steps is
        included in the returned results. This "thinning" comes at a cost of
        reduced statistical power, while reducing memory requirements and
        autocorrelation. For more discussion see [1].
        Default value: 0 (i.e., no subsampling).
      seed: Python integer to seed the random number generator.
      current_target_log_prob: (Optional) `Tensor` representing the value of
        `target_log_prob_fn` at the `current_state`. The only reason to specify
        this argument is to reduce TF graph size.
        Default value: `None` (i.e., compute as needed).
      current_grads_target_log_prob: (Optional) Python list of `Tensor`s
        representing gradient of `target_log_prob` at the `current_state` and wrt
        the `current_state`. Must have same shape as `current_state`. The only
        reason to specify this argument is to reduce TF graph size.
        Default value: `None` (i.e., compute as needed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "hmc_sample_chain").
  
    Returns:
      accepted_states: Tensor or Python list of `Tensor`s representing the
        state(s) of the Markov chain(s) at each result step. Has same shape as
        input `current_state` but with a prepended `num_results`-size dimension.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """
    with ops.name_scope(
        name, "hmc_sample_chain",
        [num_results, current_state, step_size, num_leapfrog_steps,
         num_burnin_steps, num_steps_between_results, seed,
         current_target_log_prob, current_grads_target_log_prob]):
      with ops.name_scope("initialize"):
        [
            current_state,
            step_size,
            current_target_log_prob,
            current_grads_target_log_prob,
        ] = _prepare_args(
            target_log_prob_fn,
            current_state,
            step_size,
            current_target_log_prob,
            current_grads_target_log_prob)
        num_results = ops.convert_to_tensor(
            num_results,
            dtype=dtypes.int32,
            name="num_results")
        num_leapfrog_steps = ops.convert_to_tensor(
            num_leapfrog_steps,
            dtype=dtypes.int32,
            name="num_leapfrog_steps")
        num_burnin_steps = ops.convert_to_tensor(
            num_burnin_steps,
            dtype=dtypes.int32,
            name="num_burnin_steps")
        num_steps_between_results = ops.convert_to_tensor(
            num_steps_between_results,
            dtype=dtypes.int32,
            name="num_steps_between_results")
  
      def _run_chain(num_steps, current_state, kernel_results):
        """Runs the chain(s) for `num_steps`."""
        def _loop_body(iter_, current_state, kernel_results):
          return [iter_ + 1] + list(kernel(
              target_log_prob_fn,
              current_state,
              step_size,
              num_leapfrog_steps,
              seed,
              kernel_results.current_target_log_prob,
              kernel_results.current_grads_target_log_prob))
        while_loop_kwargs = dict(
            cond=lambda iter_, *args: iter_ < num_steps,
            body=_loop_body,
            loop_vars=[
                np.int32(0),
                current_state,
                kernel_results,
            ],
        )
        if seed is not None:
          while_loop_kwargs["parallel_iterations"] = 1
        return control_flow_ops.while_loop(
            **while_loop_kwargs)[1:]  # Lop-off "iter_".
  
      def _scan_body(args_list, iter_):
        """Closure which implements `tf.scan` body."""
        current_state, kernel_results = args_list
        return _run_chain(
            1 + array_ops.where(math_ops.equal(iter_, 0),
                                num_burnin_steps,
                                num_steps_between_results),
            current_state,
            kernel_results)
  
      scan_kwargs = dict(
          fn=_scan_body,
          elems=math_ops.range(num_results),  # iter_: used to choose burnin.
          initializer=[
              current_state,
              _make_dummy_kernel_results(
                  current_state,
                  current_target_log_prob,
                  current_grads_target_log_prob),
          ])
      if seed is not None:
        scan_kwargs["parallel_iterations"] = 1
      return functional_ops.scan(**scan_kwargs)
  
  
  def sample_annealed_importance_chain(
      proposal_log_prob_fn,
      num_steps,
      target_log_prob_fn,
      current_state,
      step_size,
      num_leapfrog_steps,
      seed=None,
      name=None):
    """Runs annealed importance sampling (AIS) to estimate normalizing constants.
  
    This function uses Hamiltonian Monte Carlo to sample from a series of
    distributions that slowly interpolates between an initial "proposal"
    distribution:
  
    `exp(proposal_log_prob_fn(x) - proposal_log_normalizer)`
  
    and the target distribution:
  
    `exp(target_log_prob_fn(x) - target_log_normalizer)`,
  
    accumulating importance weights along the way. The product of these
    importance weights gives an unbiased estimate of the ratio of the
    normalizing constants of the initial distribution and the target
    distribution:
  
    `E[exp(ais_weights)] = exp(target_log_normalizer - proposal_log_normalizer)`.
  
    Note: `proposal_log_prob_fn` and `target_log_prob_fn` are called exactly three
    times (although this may be reduced to two times, in the future).
  
    #### Examples:
  
    ##### Estimate the normalizing constant of a log-gamma distribution.
  
    ```python
    tfd = tf.contrib.distributions
  
    # Run 100 AIS chains in parallel
    num_chains = 100
    dims = 20
    dtype = np.float32
  
    proposal = tfd.MultivatiateNormalDiag(
       loc=tf.zeros([dims], dtype=dtype))
  
    target = tfd.TransformedDistribution(
      distribution=tfd.Gamma(concentration=dtype(2),
                             rate=dtype(3)),
      bijector=tfd.bijectors.Invert(tfd.bijectors.Exp()),
      event_shape=[dims])
  
    chains_state, ais_weights, kernels_results = (
        hmc.sample_annealed_importance_chain(
            proposal_log_prob_fn=proposal.log_prob,
            num_steps=1000,
            target_log_prob_fn=target.log_prob,
            step_size=0.2,
            current_state=proposal.sample(num_chains),
            num_leapfrog_steps=2))
  
    log_estimated_normalizer = (tf.reduce_logsumexp(ais_weights)
                                - np.log(num_chains))
    log_true_normalizer = tf.lgamma(2.) - 2. * tf.log(3.)
    ```
  
    ##### Estimate marginal likelihood of a Bayesian regression model.
  
    ```python
    tfd = tf.contrib.distributions
  
    def make_prior(dims, dtype):
      return tfd.MultivariateNormalDiag(
          loc=tf.zeros(dims, dtype))
  
    def make_likelihood(weights, x):
      return tfd.MultivariateNormalDiag(
          loc=tf.tensordot(weights, x, axes=[[0], [-1]]))
  
    # Run 100 AIS chains in parallel
    num_chains = 100
    dims = 10
    dtype = np.float32
  
    # Make training data.
    x = np.random.randn(num_chains, dims).astype(dtype)
    true_weights = np.random.randn(dims).astype(dtype)
    y = np.dot(x, true_weights) + np.random.randn(num_chains)
  
    # Setup model.
    prior = make_prior(dims, dtype)
    def target_log_prob_fn(weights):
      return prior.log_prob(weights) + make_likelihood(weights, x).log_prob(y)
  
    proposal = tfd.MultivariateNormalDiag(
        loc=tf.zeros(dims, dtype))
  
    weight_samples, ais_weights, kernel_results = (
        hmc.sample_annealed_importance_chain(
          num_steps=1000,
          proposal_log_prob_fn=proposal.log_prob,
          target_log_prob_fn=target_log_prob_fn
          current_state=tf.zeros([num_chains, dims], dtype),
          step_size=0.1,
          num_leapfrog_steps=2))
    log_normalizer_estimate = (tf.reduce_logsumexp(ais_weights)
                               - np.log(num_chains))
    ```
  
    Args:
      proposal_log_prob_fn: Python callable that returns the log density of the
        initial distribution.
      num_steps: Integer number of Markov chain updates to run. More
        iterations means more expense, but smoother annealing between q
        and p, which in turn means exponentially lower variance for the
        normalizing constant estimator.
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step size
        for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but too-large
        step sizes make rejection exponentially more likely. When possible, it's
        often helpful to match per-variable step sizes to the standard deviations
        of the target distribution in each variable.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to `step_size *
        num_leapfrog_steps`.
      seed: Python integer to seed the random number generator.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "hmc_sample_annealed_importance_chain").
  
    Returns:
      accepted_state: `Tensor` or Python list of `Tensor`s representing the
        state(s) of the Markov chain(s) at the final iteration. Has same shape as
        input `current_state`.
      ais_weights: Tensor with the estimated weight(s). Has shape matching
        `target_log_prob_fn(current_state)`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """
    def make_convex_combined_log_prob_fn(iter_):
      def _fn(*args):
        p = proposal_log_prob_fn(*args)
        t = target_log_prob_fn(*args)
        dtype = p.dtype.base_dtype
        beta = (math_ops.cast(iter_ + 1, dtype)
                / math_ops.cast(num_steps, dtype))
        return (1. - beta) * p + beta * t
      return _fn
  
    with ops.name_scope(
        name, "hmc_sample_annealed_importance_chain",
        [num_steps, current_state, step_size, num_leapfrog_steps, seed]):
      with ops.name_scope("initialize"):
        [
            current_state,
            step_size,
            current_log_prob,
            current_grads_log_prob,
        ] = _prepare_args(
            make_convex_combined_log_prob_fn(iter_=0),
            current_state,
            step_size,
            description="convex_combined_log_prob")
        num_steps = ops.convert_to_tensor(
            num_steps,
            dtype=dtypes.int32,
            name="num_steps")
        num_leapfrog_steps = ops.convert_to_tensor(
            num_leapfrog_steps,
            dtype=dtypes.int32,
            name="num_leapfrog_steps")
      def _loop_body(iter_, ais_weights, current_state, kernel_results):
        """Closure which implements `tf.while_loop` body."""
        current_state_parts = (list(current_state)
                               if _is_list_like(current_state)
                               else [current_state])
        # TODO(b/72994218): Consider refactoring things to avoid this unecessary
        # call.
        ais_weights += ((target_log_prob_fn(*current_state_parts)
                         - proposal_log_prob_fn(*current_state_parts))
                        / math_ops.cast(num_steps, ais_weights.dtype))
        return [iter_ + 1, ais_weights] + list(kernel(
            make_convex_combined_log_prob_fn(iter_),
            current_state,
            step_size,
            num_leapfrog_steps,
            seed,
            kernel_results.current_target_log_prob,
            kernel_results.current_grads_target_log_prob))
  
      while_loop_kwargs = dict(
          cond=lambda iter_, *args: iter_ < num_steps,
          body=_loop_body,
          loop_vars=[
              np.int32(0),  # iter_
              array_ops.zeros_like(current_log_prob),  # ais_weights
              current_state,
              _make_dummy_kernel_results(current_state,
                                         current_log_prob,
                                         current_grads_log_prob),
          ])
      if seed is not None:
        while_loop_kwargs["parallel_iterations"] = 1
  
      [ais_weights, current_state, kernel_results] = control_flow_ops.while_loop(
          **while_loop_kwargs)[1:]  # Lop-off "iter_".
  
      return [current_state, ais_weights, kernel_results]
  
  
  def kernel(target_log_prob_fn,
             current_state,
             step_size,
             num_leapfrog_steps,
             seed=None,
             current_target_log_prob=None,
             current_grads_target_log_prob=None,
             name=None):
    """Runs one iteration of Hamiltonian Monte Carlo.
  
    Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC)
    algorithm that takes a series of gradient-informed steps to produce
    a Metropolis proposal. This function applies one step of HMC to
    randomly update the variable `x`.
  
    This function can update multiple chains in parallel. It assumes that all
    leftmost dimensions of `current_state` index independent chain states (and are
    therefore updated independently). The output of `target_log_prob_fn()` should
    sum log-probabilities across all event dimensions. Slices along the rightmost
    dimensions may have different target distributions; for example,
    `current_state[0, :]` could have a different target distribution from
    `current_state[1, :]`. This is up to `target_log_prob_fn()`. (The number of
    independent chains is `tf.size(target_log_prob_fn(*current_state))`.)
  
    #### Examples:
  
    ##### Simple chain with warm-up.
  
    ```python
    tfd = tf.contrib.distributions
  
    # Tuning acceptance rates:
    dtype = np.float32
    target_accept_rate = 0.631
    num_warmup_iter = 500
    num_chain_iter = 500
  
    x = tf.get_variable(name="x", initializer=dtype(1))
    step_size = tf.get_variable(name="step_size", initializer=dtype(1))
  
    target = tfd.Normal(loc=dtype(0), scale=dtype(1))
  
    new_x, other_results = hmc.kernel(
        target_log_prob_fn=target.log_prob,
        current_state=x,
        step_size=step_size,
        num_leapfrog_steps=3)[:4]
  
    x_update = x.assign(new_x)
  
    step_size_update = step_size.assign_add(
        step_size * tf.where(
          other_results.acceptance_probs > target_accept_rate,
          0.01, -0.01))
  
    warmup = tf.group([x_update, step_size_update])
  
    tf.global_variables_initializer().run()
  
    sess.graph.finalize()  # No more graph building.
  
    # Warm up the sampler and adapt the step size
    for _ in xrange(num_warmup_iter):
      sess.run(warmup)
  
    # Collect samples without adapting step size
    samples = np.zeros([num_chain_iter])
    for i in xrange(num_chain_iter):
      _, x_, target_log_prob_, grad_ = sess.run([
          x_update,
          x,
          other_results.target_log_prob,
          other_results.grads_target_log_prob])
      samples[i] = x_
  
    print(samples.mean(), samples.std())
    ```
  
    ##### Sample from more complicated posterior.
  
    I.e.,
  
    ```none
      W ~ MVN(loc=0, scale=sigma * eye(dims))
      for i=1...num_samples:
          X[i] ~ MVN(loc=0, scale=eye(dims))
        eps[i] ~ Normal(loc=0, scale=1)
          Y[i] = X[i].T * W + eps[i]
    ```
  
    ```python
    tfd = tf.contrib.distributions
  
    def make_training_data(num_samples, dims, sigma):
      dt = np.asarray(sigma).dtype
      zeros = tf.zeros(dims, dtype=dt)
      x = tfd.MultivariateNormalDiag(
          loc=zeros).sample(num_samples, seed=1)
      w = tfd.MultivariateNormalDiag(
          loc=zeros,
          scale_identity_multiplier=sigma).sample(seed=2)
      noise = tfd.Normal(
          loc=dt(0),
          scale=dt(1)).sample(num_samples, seed=3)
      y = tf.tensordot(x, w, axes=[[1], [0]]) + noise
      return y, x, w
  
    def make_prior(sigma, dims):
      # p(w | sigma)
      return tfd.MultivariateNormalDiag(
          loc=tf.zeros([dims], dtype=sigma.dtype),
          scale_identity_multiplier=sigma)
  
    def make_likelihood(x, w):
      # p(y | x, w)
      return tfd.MultivariateNormalDiag(
          loc=tf.tensordot(x, w, axes=[[1], [0]]))
  
    # Setup assumptions.
    dtype = np.float32
    num_samples = 150
    dims = 10
    num_iters = int(5e3)
  
    true_sigma = dtype(0.5)
    y, x, true_weights = make_training_data(num_samples, dims, true_sigma)
  
    # Estimate of `log(true_sigma)`.
    log_sigma = tf.get_variable(name="log_sigma", initializer=dtype(0))
    sigma = tf.exp(log_sigma)
  
    # State of the Markov chain.
    weights = tf.get_variable(
        name="weights",
        initializer=np.random.randn(dims).astype(dtype))
  
    prior = make_prior(sigma, dims)
  
    def joint_log_prob_fn(w):
      # f(w) = log p(w, y | x)
      return prior.log_prob(w) + make_likelihood(x, w).log_prob(y)
  
    weights_update = weights.assign(
        hmc.kernel(target_log_prob_fn=joint_log_prob,
                   current_state=weights,
                   step_size=0.1,
                   num_leapfrog_steps=5)[0])
  
    with tf.control_dependencies([weights_update]):
      loss = -prior.log_prob(weights)
  
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)
    log_sigma_update = optimizer.minimize(loss, var_list=[log_sigma])
  
    sess.graph.finalize()  # No more graph building.
  
    tf.global_variables_initializer().run()
  
    sigma_history = np.zeros(num_iters, dtype)
    weights_history = np.zeros([num_iters, dims], dtype)
  
    for i in xrange(num_iters):
      _, sigma_, weights_, _ = sess.run([log_sigma_update, sigma, weights])
      weights_history[i, :] = weights_
      sigma_history[i] = sigma_
  
    true_weights_ = sess.run(true_weights)
  
    # Should converge to something close to true_sigma.
    plt.plot(sigma_history);
    plt.ylabel("sigma");
    plt.xlabel("iteration");
    ```
  
    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `current_state` (or `*current_state` if it's a list) and returns its
        (possibly unnormalized) log-density under the target distribution.
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      step_size: `Tensor` or Python `list` of `Tensor`s representing the step size
        for the leapfrog integrator. Must broadcast with the shape of
        `current_state`. Larger step sizes lead to faster progress, but too-large
        step sizes make rejection exponentially more likely. When possible, it's
        often helpful to match per-variable step sizes to the standard deviations
        of the target distribution in each variable.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to `step_size *
        num_leapfrog_steps`.
      seed: Python integer to seed the random number generator.
      current_target_log_prob: (Optional) `Tensor` representing the value of
        `target_log_prob_fn` at the `current_state`. The only reason to
        specify this argument is to reduce TF graph size.
        Default value: `None` (i.e., compute as needed).
      current_grads_target_log_prob: (Optional) Python list of `Tensor`s
        representing gradient of `current_target_log_prob` at the `current_state`
        and wrt the `current_state`. Must have same shape as `current_state`. The
        only reason to specify this argument is to reduce TF graph size.
        Default value: `None` (i.e., compute as needed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "hmc_kernel").
  
    Returns:
      accepted_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) at each result step. Has same shape as
        `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
  
    Raises:
      ValueError: if there isn't one `step_size` or a list with same length as
        `current_state`.
    """
    with ops.name_scope(
        name, "hmc_kernel",
        [current_state, step_size, num_leapfrog_steps, seed,
         current_target_log_prob, current_grads_target_log_prob]):
      with ops.name_scope("initialize"):
        [current_state_parts, step_sizes, current_target_log_prob,
         current_grads_target_log_prob] = _prepare_args(
             target_log_prob_fn, current_state, step_size,
             current_target_log_prob, current_grads_target_log_prob,
             maybe_expand=True)
        independent_chain_ndims = distributions_util.prefer_static_rank(
            current_target_log_prob)
        current_momentums = []
        for s in current_state_parts:
          current_momentums.append(random_ops.random_normal(
              shape=array_ops.shape(s),
              dtype=s.dtype.base_dtype,
              seed=seed))
          seed = distributions_util.gen_new_seed(
              seed, salt="hmc_kernel_momentums")
  
        num_leapfrog_steps = ops.convert_to_tensor(
            num_leapfrog_steps,
            dtype=dtypes.int32,
            name="num_leapfrog_steps")
      [
          proposed_momentums,
          proposed_state_parts,
          proposed_target_log_prob,
          proposed_grads_target_log_prob,
      ] = _leapfrog_integrator(current_momentums,
                               target_log_prob_fn,
                               current_state_parts,
                               step_sizes,
                               num_leapfrog_steps,
                               current_target_log_prob,
                               current_grads_target_log_prob)
  
      energy_change = _compute_energy_change(current_target_log_prob,
                                             current_momentums,
                                             proposed_target_log_prob,
                                             proposed_momentums,
                                             independent_chain_ndims)
  
      # u < exp(min(-energy, 0)),  where u~Uniform[0,1)
      # ==> -log(u) >= max(e, 0)
      # ==> -log(u) >= e
      # (Perhaps surprisingly, we don't have a better way to obtain a random
      # uniform from positive reals, i.e., `tf.random_uniform(minval=0,
      # maxval=np.inf)` won't work.)
      random_uniform = random_ops.random_uniform(
          shape=array_ops.shape(energy_change),
          dtype=energy_change.dtype,
          seed=seed)
      random_positive = -math_ops.log(random_uniform)
      is_accepted = random_positive >= energy_change
  
      accepted_target_log_prob = array_ops.where(is_accepted,
                                                 proposed_target_log_prob,
                                                 current_target_log_prob)
  
      accepted_state_parts = [_choose(is_accepted,
                                      proposed_state_part,
                                      current_state_part,
                                      independent_chain_ndims)
                              for current_state_part, proposed_state_part
                              in zip(current_state_parts, proposed_state_parts)]
  
      accepted_grads_target_log_prob = [
          _choose(is_accepted,
                  proposed_grad,
                  grad,
                  independent_chain_ndims)
          for proposed_grad, grad
          in zip(proposed_grads_target_log_prob, current_grads_target_log_prob)]
  
      maybe_flatten = lambda x: x if _is_list_like(current_state) else x[0]
      return [
          maybe_flatten(accepted_state_parts),
          KernelResults(
              acceptance_probs=math_ops.exp(math_ops.minimum(-energy_change, 0.)),
              current_grads_target_log_prob=accepted_grads_target_log_prob,
              current_target_log_prob=accepted_target_log_prob,
              energy_change=energy_change,
              is_accepted=is_accepted,
              proposed_grads_target_log_prob=proposed_grads_target_log_prob,
              proposed_state=maybe_flatten(proposed_state_parts),
              proposed_target_log_prob=proposed_target_log_prob,
              random_positive=random_positive,
          ),
      ]
  
  
  def _leapfrog_integrator(current_momentums,
                           target_log_prob_fn,
                           current_state_parts,
                           step_sizes,
                           num_leapfrog_steps,
                           current_target_log_prob=None,
                           current_grads_target_log_prob=None,
                           name=None):
    """Applies `num_leapfrog_steps` of the leapfrog integrator.
  
    Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`.
  
    #### Examples:
  
    ##### Simple quadratic potential.
  
    ```python
    tfd = tf.contrib.distributions
  
    dims = 10
    num_iter = int(1e3)
    dtype = np.float32
  
    position = tf.placeholder(np.float32)
    momentum = tf.placeholder(np.float32)
  
    [
        new_momentums,
        new_positions,
    ] = hmc._leapfrog_integrator(
        current_momentums=[momentum],
        target_log_prob_fn=tfd.MultivariateNormalDiag(
            loc=tf.zeros(dims, dtype)).log_prob,
        current_state_parts=[position],
        step_sizes=0.1,
        num_leapfrog_steps=3)[:2]
  
    sess.graph.finalize()  # No more graph building.
  
    momentum_ = np.random.randn(dims).astype(dtype)
    position_ = np.random.randn(dims).astype(dtype)
  
    positions = np.zeros([num_iter, dims], dtype)
    for i in xrange(num_iter):
      position_, momentum_ = sess.run(
          [new_momentums[0], new_position[0]],
          feed_dict={position: position_, momentum: momentum_})
      positions[i] = position_
  
    plt.plot(positions[:, 0]);  # Sinusoidal.
    ```
  
    Args:
      current_momentums: Tensor containing the value(s) of the momentum
        variable(s) to update.
      target_log_prob_fn: Python callable which takes an argument like
        `*current_state_parts` and returns its (possibly unnormalized) log-density
        under the target distribution.
      current_state_parts: Python `list` of `Tensor`s representing the current
        state(s) of the Markov chain(s). The first `independent_chain_ndims` of
        the `Tensor`(s) index different chains.
      step_sizes: Python `list` of `Tensor`s representing the step size for the
        leapfrog integrator. Must broadcast with the shape of
        `current_state_parts`.  Larger step sizes lead to faster progress, but
        too-large step sizes make rejection exponentially more likely. When
        possible, it's often helpful to match per-variable step sizes to the
        standard deviations of the target distribution in each variable.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to `step_size *
        num_leapfrog_steps`.
      current_target_log_prob: (Optional) `Tensor` representing the value of
        `target_log_prob_fn(*current_state_parts)`. The only reason to specify
        this argument is to reduce TF graph size.
        Default value: `None` (i.e., compute as needed).
      current_grads_target_log_prob: (Optional) Python list of `Tensor`s
        representing gradient of `target_log_prob_fn(*current_state_parts`) wrt
        `current_state_parts`. Must have same shape as `current_state_parts`. The
        only reason to specify this argument is to reduce TF graph size.
        Default value: `None` (i.e., compute as needed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., "hmc_leapfrog_integrator").
  
    Returns:
      proposed_momentums: Updated value of the momentum.
      proposed_state_parts: Tensor or Python list of `Tensor`s representing the
        state(s) of the Markov chain(s) at each result step. Has same shape as
        input `current_state_parts`.
      proposed_target_log_prob: `Tensor` representing the value of
        `target_log_prob_fn` at `accepted_state`.
      proposed_grads_target_log_prob: Gradient of `proposed_target_log_prob` wrt
        `accepted_state`.
  
    Raises:
      ValueError: if `len(momentums) != len(state_parts)`.
      ValueError: if `len(state_parts) != len(step_sizes)`.
      ValueError: if `len(state_parts) != len(grads_target_log_prob)`.
      TypeError: if `not target_log_prob.dtype.is_floating`.
    """
    def _loop_body(step,
                   current_momentums,
                   current_state_parts,
                   ignore_current_target_log_prob,  # pylint: disable=unused-argument
                   current_grads_target_log_prob):
      return [step + 1] + list(_leapfrog_step(current_momentums,
                                              target_log_prob_fn,
                                              current_state_parts,
                                              step_sizes,
                                              current_grads_target_log_prob))
  
    with ops.name_scope(
        name, "hmc_leapfrog_integrator",
        [current_momentums, current_state_parts, step_sizes, num_leapfrog_steps,
         current_target_log_prob, current_grads_target_log_prob]):
      if len(current_momentums) != len(current_state_parts):
        raise ValueError("`momentums` must be in one-to-one correspondence "
                         "with `state_parts`")
      num_leapfrog_steps = ops.convert_to_tensor(num_leapfrog_steps,
                                                 name="num_leapfrog_steps")
      current_target_log_prob, current_grads_target_log_prob = (
          _maybe_call_fn_and_grads(
              target_log_prob_fn,
              current_state_parts,
              current_target_log_prob,
              current_grads_target_log_prob))
      return control_flow_ops.while_loop(
          cond=lambda iter_, *args: iter_ < num_leapfrog_steps,
          body=_loop_body,
          loop_vars=[
              np.int32(0),  # iter_
              current_momentums,
              current_state_parts,
              current_target_log_prob,
              current_grads_target_log_prob,
          ],
          back_prop=False)[1:]  # Lop-off "iter_".
  
  
  def _leapfrog_step(current_momentums,
                     target_log_prob_fn,
                     current_state_parts,
                     step_sizes,
                     current_grads_target_log_prob,
                     name=None):
    """Applies one step of the leapfrog integrator."""
    with ops.name_scope(
        name, "_leapfrog_step",
        [current_momentums, current_state_parts, step_sizes,
         current_grads_target_log_prob]):
      proposed_momentums = [m + 0.5 * ss * g for m, ss, g
                            in zip(current_momentums,
                                   step_sizes,
                                   current_grads_target_log_prob)]
      proposed_state_parts = [x + ss * m for x, ss, m
                              in zip(current_state_parts,
                                     step_sizes,
                                     proposed_momentums)]
      proposed_target_log_prob = target_log_prob_fn(*proposed_state_parts)
      if not proposed_target_log_prob.dtype.is_floating:
        raise TypeError("`target_log_prob_fn` must produce a `Tensor` "
                        "with `float` `dtype`.")
      proposed_grads_target_log_prob = gradients_ops.gradients(
          proposed_target_log_prob, proposed_state_parts)
      if any(g is None for g in proposed_grads_target_log_prob):
        raise ValueError(
            "Encountered `None` gradient. Does your target `target_log_prob_fn` "
            "access all `tf.Variable`s via `tf.get_variable`?\n"
            "  current_state_parts: {}\n"
            "  proposed_state_parts: {}\n"
            "  proposed_grads_target_log_prob: {}".format(
                current_state_parts,
                proposed_state_parts,
                proposed_grads_target_log_prob))
      proposed_momentums = [m + 0.5 * ss * g for m, ss, g
                            in zip(proposed_momentums,
                                   step_sizes,
                                   proposed_grads_target_log_prob)]
      return [
          proposed_momentums,
          proposed_state_parts,
          proposed_target_log_prob,
          proposed_grads_target_log_prob,
      ]
  
  
  def _compute_energy_change(current_target_log_prob,
                             current_momentums,
                             proposed_target_log_prob,
                             proposed_momentums,
                             independent_chain_ndims,
                             name=None):
    """Helper to `kernel` which computes the energy change."""
    with ops.name_scope(
        name, "compute_energy_change",
        ([current_target_log_prob, proposed_target_log_prob,
          independent_chain_ndims] +
         current_momentums + proposed_momentums)):
      # Abbreviate lk0=log_kinetic_energy and lk1=proposed_log_kinetic_energy
      # since they're a mouthful and lets us inline more.
      lk0, lk1 = [], []
      for current_momentum, proposed_momentum in zip(current_momentums,
                                                     proposed_momentums):
        axis = math_ops.range(independent_chain_ndims,
                              array_ops.rank(current_momentum))
        lk0.append(_log_sum_sq(current_momentum, axis))
        lk1.append(_log_sum_sq(proposed_momentum, axis))
  
      lk0 = -np.log(2.) + math_ops.reduce_logsumexp(array_ops.stack(lk0, axis=-1),
                                                    axis=-1)
      lk1 = -np.log(2.) + math_ops.reduce_logsumexp(array_ops.stack(lk1, axis=-1),
                                                    axis=-1)
      lp0 = -current_target_log_prob   # log_potential
      lp1 = -proposed_target_log_prob  # proposed_log_potential
      x = array_ops.stack([lp1, math_ops.exp(lk1), -lp0, -math_ops.exp(lk0)],
                          axis=-1)
  
      # The sum is NaN if any element is NaN or we see both +Inf and -Inf.
      # Thus we will replace such rows with infinite energy change which implies
      # rejection. Recall that float-comparisons with NaN are always False.
      is_sum_determinate = (
          math_ops.reduce_all(math_ops.is_finite(x) | (x >= 0.), axis=-1) &
          math_ops.reduce_all(math_ops.is_finite(x) | (x <= 0.), axis=-1))
      is_sum_determinate = array_ops.tile(
          is_sum_determinate[..., array_ops.newaxis],
          multiples=array_ops.concat([
              array_ops.ones(array_ops.rank(is_sum_determinate),
                             dtype=dtypes.int32),
              [4],
          ], axis=0))
      x = array_ops.where(is_sum_determinate,
                          x,
                          array_ops.fill(array_ops.shape(x),
                                         value=x.dtype.as_numpy_dtype(np.inf)))
  
      return math_ops.reduce_sum(x, axis=-1)
  
  
  def _choose(is_accepted,
              accepted,
              rejected,
              independent_chain_ndims,
              name=None):
    """Helper to `kernel` which expand_dims `is_accepted` to apply tf.where."""
    def _expand_is_accepted_like(x):
      with ops.name_scope("_choose"):
        expand_shape = array_ops.concat([
            array_ops.shape(is_accepted),
            array_ops.ones([array_ops.rank(x) - array_ops.rank(is_accepted)],
                           dtype=dtypes.int32),
        ], axis=0)
        multiples = array_ops.concat([
            array_ops.ones([array_ops.rank(is_accepted)], dtype=dtypes.int32),
            array_ops.shape(x)[independent_chain_ndims:],
        ], axis=0)
        m = array_ops.tile(array_ops.reshape(is_accepted, expand_shape),
                           multiples)
        m.set_shape(x.shape)
        return m
    with ops.name_scope(name, "_choose", values=[
        is_accepted, accepted, rejected, independent_chain_ndims]):
      return array_ops.where(_expand_is_accepted_like(accepted),
                             accepted,
                             rejected)
  
  
  def _maybe_call_fn_and_grads(fn,
                               fn_arg_list,
                               fn_result=None,
                               grads_fn_result=None,
                               description="target_log_prob"):
    """Helper which computes `fn_result` and `grads` if needed."""
    fn_arg_list = (list(fn_arg_list) if _is_list_like(fn_arg_list)
                   else [fn_arg_list])
    if fn_result is None:
      fn_result = fn(*fn_arg_list)
    if not fn_result.dtype.is_floating:
      raise TypeError("`{}` must be a `Tensor` with `float` `dtype`.".format(
          description))
    if grads_fn_result is None:
      grads_fn_result = gradients_ops.gradients(
          fn_result, fn_arg_list)
    if len(fn_arg_list) != len(grads_fn_result):
      raise ValueError("`{}` must be in one-to-one correspondence with "
                       "`grads_{}`".format(*[description]*2))
    if any(g is None for g in grads_fn_result):
      raise ValueError("Encountered `None` gradient.")
    return fn_result, grads_fn_result
  
  
  def _prepare_args(target_log_prob_fn, state, step_size,
                    target_log_prob=None, grads_target_log_prob=None,
                    maybe_expand=False, description="target_log_prob"):
    """Helper which processes input args to meet list-like assumptions."""
    state_parts = list(state) if _is_list_like(state) else [state]
    state_parts = [ops.convert_to_tensor(s, name="state")
                   for s in state_parts]
    target_log_prob, grads_target_log_prob = _maybe_call_fn_and_grads(
        target_log_prob_fn,
        state_parts,
        target_log_prob,
        grads_target_log_prob,
        description)
    step_sizes = list(step_size) if _is_list_like(step_size) else [step_size]
    step_sizes = [
        ops.convert_to_tensor(
            s, name="step_size", dtype=target_log_prob.dtype)
        for s in step_sizes]
    if len(step_sizes) == 1:
      step_sizes *= len(state_parts)
    if len(state_parts) != len(step_sizes):
      raise ValueError("There should be exactly one `step_size` or it should "
                       "have same length as `current_state`.")
    maybe_flatten = lambda x: x if maybe_expand or _is_list_like(state) else x[0]
    return [
        maybe_flatten(state_parts),
        maybe_flatten(step_sizes),
        target_log_prob,
        grads_target_log_prob,
    ]
  
  
  def _is_list_like(x):
    """Helper which returns `True` if input is `list`-like."""
    return isinstance(x, (tuple, list))
  
  
  def _log_sum_sq(x, axis=None):
    """Computes log(sum(x**2))."""
    return math_ops.reduce_logsumexp(2. * math_ops.log(math_ops.abs(x)), axis)
  
apis: tensorflow.python.ops.array_ops.sequence_mask tensorflow.python.ops.nn_ops.bias_add_grad tensorflow.python.ops.array_ops.expand_dims tensorflow.python.ops.array_ops.reshape tensorflow.python.platform.resource_loader.get_path_to_datafile tensorflow.python.ops.math_ops.matmul tensorflow.python.ops.array_ops.slice tensorflow.python.ops.math_ops.range tensorflow.python.ops.init_ops.constant_initializer tensorflow.python.ops.array_ops.unstack tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.array_ops.shape tensorflow.contrib.rnn.ops.gen_lstm_ops.block_lstm tensorflow.python.ops.array_ops.constant tensorflow.python.framework.ops.convert_to_tensor tensorflow.contrib.rnn.ops.gen_lstm_ops.lstm_block_cell tensorflow.python.layers.base.InputSpec tensorflow.python.framework.ops.RegisterGradient tensorflow.python.ops.math_ops.reduce_max tensorflow.python.ops.array_ops.stack tensorflow.python.ops.array_ops.concat tensorflow.python.ops.rnn_cell_impl.LSTMStateTuple tensorflow.python.ops.math_ops.to_int64
apis: tensorflow.python.ops.math_ops.equal tensorflow.python.ops.distributions.util.prefer_static_rank tensorflow.python.ops.gradients_impl.gradients numpy.log tensorflow.python.ops.math_ops.is_finite tensorflow.python.ops.math_ops.range tensorflow.python.ops.array_ops.rank tensorflow.python.ops.math_ops.minimum tensorflow.python.ops.math_ops.reduce_sum tensorflow.python.ops.array_ops.shape tensorflow.python.ops.control_flow_ops.while_loop numpy.int32 tensorflow.python.ops.math_ops.cast tensorflow.python.ops.functional_ops.scan tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.ops.math_ops.abs tensorflow.python.ops.array_ops.zeros_like tensorflow.python.ops.math_ops.log tensorflow.python.framework.ops.name_scope tensorflow.python.ops.array_ops.ones_like tensorflow.python.ops.array_ops.stack tensorflow.python.ops.array_ops.where tensorflow.python.ops.array_ops.reshape tensorflow.python.ops.distributions.util.gen_new_seed tensorflow.python.ops.math_ops.exp