def log_joint_fn(*param_vals):
        """Generated log-density function."""

        # Sum the log_prob values from parameter priors.
        param_lp = sum([
            param.prior.log_prob(param_val)
            for (param, param_val) in zip(self.parameters, param_vals)
        ])

        # Build a linear Gaussian state space model and evaluate the marginal
        # log_prob on observations.
        lgssm = self.make_state_space_model(
            param_vals=param_vals, num_timesteps=num_timesteps)
        observation_lp = lgssm.log_prob(observed_time_series)

        # Sum over likelihoods from iid observations. Without this sum,
        # adding `param_lp + observation_lp` would broadcast the param priors
        # over the sample shape, which incorrectly multi-counts the param
        # priors.
        sample_ndims = tf.maximum(0,
                                  tf.rank(observation_lp) - tf.rank(param_lp))
        observation_lp = tf.reduce_sum(
            observation_lp, axis=tf.range(sample_ndims))

        return param_lp + observation_lp

  def testDenseShape(self):
    with self.test_session():
      t_value = [[0, 42], [24, 0]]
      self.assertAllEqual((2, 2), tf.shape(t_value).eval())
      self.assertEqual(4, tf.size(t_value).eval())
      self.assertEqual(2, tf.rank(t_value).eval())

      t = tf.constant(t_value)
      self.assertAllEqual((2, 2), tf.shape(t).eval())
      self.assertEqual(4, tf.size(t).eval())
      self.assertEqual(2, tf.rank(t).eval())

    def testSparseShape(self):
        with self.test_session():
            sp_value = tf.SparseTensorValue(indices=((0, 1), (1, 0)), values=(42, 24), shape=(2, 2))
            self.assertAllEqual((2, 2), tf.shape(sp_value).eval())
            self.assertEqual(4, tf.size(sp_value).eval())
            self.assertEqual(2, tf.rank(sp_value).eval())

            sp = tf.SparseTensor.from_value(sp_value)
            self.assertAllEqual((2, 2), tf.shape(sp).eval())
            self.assertEqual(4, tf.size(sp).eval())
            self.assertEqual(2, tf.rank(sp).eval())

 def new_target_log_prob(*transformed_state_parts):
   """Log prob of the transformed state."""
   # TODO(b/72831017): Use `tf.identity` to disable caching (since HMC takes
   # gradient with respect to input).
   transformed_state_parts = [
       tf.identity(sp) for sp in transformed_state_parts
   ]
   tlp = target_log_prob_fn(
       *self._forward_transform(transformed_state_parts))
   event_ndims = [
       tf.rank(sp) - tf.rank(tlp) for sp in transformed_state_parts
   ]
   return tlp + self._forward_log_det_jacobian(
       transformed_state_parts=transformed_state_parts,
       event_ndims=event_ndims)

 def _pad_sample_dims(self, x):
   with tf.name_scope("pad_sample_dims", values=[x]):
     ndims = x.shape.ndims if x.shape.ndims is not None else tf.rank(x)
     shape = tf.shape(x)
     d = ndims - self._event_ndims
     x = tf.reshape(x, shape=tf.concat([shape[:d], [1], shape[d:]], axis=0))
     return x

def _transpose_batch_time(x):
    """Transpose the batch and time dimensions of a Tensor.

    Retains as much of the static shape information as possible.

    Args:
        x: A tensor of rank 2 or higher.

    Returns: x transposed along the first two dimensions.

    Raises:
        ValueError: if `x` is rank 1 or lower.
    """
    x_static_shape = x.get_shape()
    if x_static_shape.ndims is not None and x_static_shape.ndims < 2:
        raise ValueError(
            "Expected input tensor %s to have rank at least 2, but saw shape: %s" %
            (x, x_static_shape))
    x_rank = tf.rank(x)
    x_t = tf.transpose(
        x, tf.concat(
            ([1, 0], tf.range(2, x_rank)), axis=0))
    x_t.set_shape(
        tf.TensorShape([
            x_static_shape[1].value, x_static_shape[0].value
        ]).concatenate(x_static_shape[2:]))
    return x_t

def _do_maximum_mean(samples, envelope, high, name=None):
  """Common code between maximum_mean and minimum_mean."""
  with tf.name_scope(name, "do_maximum_mean", [samples, envelope, high]):
    dtype = dtype_util.common_dtype([samples, envelope, high], tf.float32)
    samples = tf.convert_to_tensor(samples, name="samples", dtype=dtype)
    envelope = tf.convert_to_tensor(envelope, name="envelope", dtype=dtype)
    high = tf.convert_to_tensor(high, name="high", dtype=dtype)
    n = tf.rank(samples)
    # Move the batch dimension of `samples` to the rightmost position,
    # where the _batch_sort_vector function wants it.
    perm = tf.concat([tf.range(1, n), [0]], axis=0)
    samples = tf.transpose(samples, perm)

    samples = _batch_sort_vector(samples)

    # The maximum mean is given by taking `envelope`-worth of
    # probability from the smallest samples and moving it to the
    # maximum value.  This amounts to:
    # - ignoring the smallest k samples, where `k/n < envelope`
    # - taking a `1/n - (envelope - k/n)` part of the index k sample
    # - taking all the other samples
    # - and adding `envelope * high` at the end.
    # The following is a vectorized and batched way of computing this.
    # `max_mean_contrib` is a mask implementing the previous.
    batch_size = tf.shape(samples)[-1]
    batch_size = tf.cast(batch_size, dtype=dtype)
    step = 1. / batch_size
    cum_steps = step * tf.range(1, batch_size + 1, dtype=dtype)
    max_mean_contrib = tf.clip_by_value(
        cum_steps - envelope[..., tf.newaxis],
        clip_value_min=0.,
        clip_value_max=step)
    return tf.reduce_sum(samples * max_mean_contrib, axis=-1) + envelope * high

  def _make_columnar(self, x):
    """Ensures non-scalar input has at least one column.

    Example:
      If `x = [1, 2, 3]` then the output is `[[1], [2], [3]]`.

      If `x = [[1, 2, 3], [4, 5, 6]]` then the output is unchanged.

      If `x = 1` then the output is unchanged.

    Args:
      x: `Tensor`.

    Returns:
      columnar_x: `Tensor` with at least two dimensions.
    """
    if x.shape.ndims is not None:
      if x.shape.ndims == 1:
        x = x[tf.newaxis, :]
      return x
    shape = tf.shape(x)
    maybe_expanded_shape = tf.concat([
        shape[:-1],
        distribution_util.pick_vector(
            tf.equal(tf.rank(x), 1), [1], np.array([], dtype=np.int32)),
        shape[-1:],
    ], 0)
    return tf.reshape(x, maybe_expanded_shape)

  def _inverse_log_det_jacobian(self, y, **kwargs):
    y = tf.convert_to_tensor(y, name="y")
    ildj = tf.cast(0., dtype=y.dtype.base_dtype)

    if not self.bijectors:
      return ildj

    event_ndims = self._maybe_get_static_event_ndims(
        self.inverse_min_event_ndims)

    if _use_static_shape(y, event_ndims):
      event_shape = y.shape[y.shape.ndims - event_ndims:]
    else:
      event_shape = tf.shape(y)[tf.rank(y) - event_ndims:]

    for b in self.bijectors:
      ildj += b.inverse_log_det_jacobian(
          y, event_ndims=event_ndims, **kwargs.get(b.name, {}))

      if _use_static_shape(y, event_ndims):
        event_shape = b.inverse_event_shape(event_shape)
        event_ndims = self._maybe_get_static_event_ndims(
            event_shape.ndims)
      else:
        event_shape = b.inverse_event_shape_tensor(event_shape)
        event_ndims = tf.size(event_shape)
        event_ndims_ = self._maybe_get_static_event_ndims(event_ndims)
        if event_ndims_ is not None:
          event_ndims = event_ndims_

      y = b.inverse(y, **kwargs.get(b.name, {}))
    return ildj

  def _forward_log_det_jacobian(self, x, **kwargs):
    x = tf.convert_to_tensor(x, name="x")

    fldj = tf.cast(0., dtype=x.dtype.base_dtype)

    if not self.bijectors:
      return fldj

    event_ndims = self._maybe_get_static_event_ndims(
        self.forward_min_event_ndims)

    if _use_static_shape(x, event_ndims):
      event_shape = x.shape[x.shape.ndims - event_ndims:]
    else:
      event_shape = tf.shape(x)[tf.rank(x) - event_ndims:]

    for b in reversed(self.bijectors):
      fldj += b.forward_log_det_jacobian(
          x, event_ndims=event_ndims, **kwargs.get(b.name, {}))
      if _use_static_shape(x, event_ndims):
        event_shape = b.forward_event_shape(event_shape)
        event_ndims = self._maybe_get_static_event_ndims(event_shape.ndims)
      else:
        event_shape = b.forward_event_shape_tensor(event_shape)
        event_ndims = tf.size(event_shape)
        event_ndims_ = self._maybe_get_static_event_ndims(event_ndims)
        if event_ndims_ is not None:
          event_ndims = event_ndims_

      x = b.forward(x, **kwargs.get(b.name, {}))

    return fldj

 def _expand_is_accepted_like(x):
   """Helper to expand `is_accepted` like the shape of some input arg."""
   with tf.name_scope('expand_is_accepted_like'):
     expand_shape = tf.concat([
         tf.shape(is_accepted),
         tf.ones([tf.rank(x) - tf.rank(is_accepted)],
                 dtype=tf.int32),
     ], axis=0)
     multiples = tf.concat([
         tf.ones([tf.rank(is_accepted)], dtype=tf.int32),
         tf.shape(x)[tf.rank(is_accepted):],
     ], axis=0)
     m = tf.tile(tf.reshape(is_accepted, expand_shape),
                 multiples)
     m.set_shape(m.shape.merge_with(x.shape))
     return m

 def _compareRank(self, x, use_gpu=False):
   np_ans = np.asarray(np.ndim(x))
   with self.test_session(use_gpu=use_gpu):
     tf_ans = tf.rank(x)
     result = tf_ans.eval()
   self.assertAllEqual(np_ans, result)
   self.assertShapeEqual(np_ans, tf_ans)

def _tensor_product(t1, t2):
  """Computes the tensor product of two tensors.

  If the rank of `t1` is `q` and the rank of `t2` is `r`, the result `z` is
  of rank `q+r` with shape `t1.shape + t2.shape`. The components of `z` are:

  ```None
    z[i1, i2, .., iq, j1, j2, .., jr] = t1[i1, .., iq] * t2[j1, .., jq]
  ```

  If both inputs are of rank 1, then the tensor product is equivalent to outer
  product of vectors.

  Note that tensor product is not commutative in general.

  Args:
    t1: A `tf.Tensor` of any dtype and non zero rank.
    t2: A `tf.Tensor` of same dtype as `t1` and non zero rank.

  Returns:
    product: A tensor with the same elements as the input `x` but with rank
      `r + n` where `r` is the rank of `x`.
  """
  t1_shape = tf.shape(t1)
  padding = tf.ones([tf.rank(t2)], dtype=t1_shape.dtype)
  padded_shape = tf.concat([t1_shape, padding], axis=0)
  t1_padded = tf.reshape(t1, padded_shape)
  return t1_padded * t2

  def _expand_sample_shape_to_vector(self, x, name):
    """Helper to `sample` which ensures input is 1D."""
    x_static_val = tf.contrib.util.constant_value(x)
    if x_static_val is None:
      prod = tf.reduce_prod(x)
    else:
      prod = np.prod(x_static_val, dtype=x.dtype.as_numpy_dtype())

    ndims = x.shape.ndims  # != sample_ndims
    if ndims is None:
      # Maybe expand_dims.
      ndims = tf.rank(x)
      expanded_shape = util.pick_vector(
          tf.equal(ndims, 0),
          np.array([1], dtype=np.int32), tf.shape(x))
      x = tf.reshape(x, expanded_shape)
    elif ndims == 0:
      # Definitely expand_dims.
      if x_static_val is not None:
        x = tf.convert_to_tensor(
            np.array([x_static_val], dtype=x.dtype.as_numpy_dtype()),
            name=name)
      else:
        x = tf.reshape(x, [1])
    elif ndims != 1:
      raise ValueError("Input is neither scalar nor vector.")

    return x, prod

 def _transpose(self, x, perm):
   sample_batch_ndims = tf.rank(x) - self.rightmost_transposed_ndims
   perm = tf.concat([
       tf.range(sample_batch_ndims),
       sample_batch_ndims + perm,
   ], axis=0)
   return tf.transpose(x, perm)

 def _compareRankSparse(self, x_np, use_gpu=False):
   np_ans = np.asarray(np.ndim(x_np))
   x_tf, unused_nnz = _sparsify(x_np)
   with self.test_session(use_gpu=use_gpu):
     tf_ans = tf.rank(x_tf)
     result = tf_ans.eval()
   self.assertAllEqual(np_ans, result)
   self.assertShapeEqual(np_ans, tf_ans)

def _squeeze(x, axis):
  """A version of squeeze that works with dynamic axis."""
  x = tf.convert_to_tensor(x, name='x')
  if axis is None:
    return tf.squeeze(x, axis=None)
  axis = tf.convert_to_tensor(axis, name='axis', dtype=tf.int32)
  axis += tf.zeros([1], dtype=axis.dtype)  # Make axis at least 1d.
  keep_axis, _ = tf.setdiff1d(tf.range(0, tf.rank(x)), axis)
  return tf.reshape(x, tf.gather(tf.shape(x), keep_axis))

 def _maybe_rotate_dims(self, x, rotate_right=False):
   """Helper which rolls left event_dims left or right event_dims right."""
   needs_rotation_const = tf.contrib.util.constant_value(self._needs_rotation)
   if needs_rotation_const is not None and not needs_rotation_const:
     return x
   ndims = tf.rank(x)
   n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims
   return tf.transpose(
       x, _concat_vectors(tf.range(n, ndims), tf.range(0, n)))

  def __init__(self,
               logits=None,
               probs=None,
               dtype=tf.int32,
               validate_args=False,
               allow_nan_stats=True,
               name="OneHotCategorical"):
    """Initialize OneHotCategorical distributions using class log-probabilities.

    Args:
      logits: An N-D `Tensor`, `N >= 1`, representing the log probabilities of a
        set of Categorical distributions. The first `N - 1` dimensions index
        into a batch of independent distributions and the last dimension
        represents a vector of logits for each class. Only one of `logits` or
        `probs` should be passed in.
      probs: An N-D `Tensor`, `N >= 1`, representing the probabilities of a set
        of Categorical distributions. The first `N - 1` dimensions index into a
        batch of independent distributions and the last dimension represents a
        vector of probabilities for each class. Only one of `logits` or `probs`
        should be passed in.
      dtype: The type of the event samples (default: int32).
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    with tf.name_scope(name, values=[logits, probs]) as name:
      self._logits, self._probs = distribution_util.get_logits_and_probs(
          name=name, logits=logits, probs=probs, validate_args=validate_args,
          multidimensional=True)

      logits_shape_static = self._logits.shape.with_rank_at_least(1)
      if logits_shape_static.ndims is not None:
        self._batch_rank = tf.convert_to_tensor(
            logits_shape_static.ndims - 1, dtype=tf.int32, name="batch_rank")
      else:
        with tf.name_scope(name="batch_rank"):
          self._batch_rank = tf.rank(self._logits) - 1

      with tf.name_scope(name="event_size"):
        self._event_size = tf.shape(self._logits)[-1]

    super(OneHotCategorical, self).__init__(
        dtype=dtype,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._logits, self._probs],
        name=name)

	def step(self, x, c=None, g=None, softmax=False):
		"""Forward step

		Args:
			x: Tensor of shape [batch_size, channels, time_length], One-hot encoded audio signal.
			c: Tensor of shape [batch_size, cin_channels, time_length], Local conditioning features.
			g: Tensor of shape [batch_size, gin_channels, 1] or Ids of shape [batch_size, 1], 
				Global conditioning features.
				Note: set hparams.use_speaker_embedding to False to disable embedding layer and 
				use extrnal One-hot encoded features.
			softmax: Boolean, Whether to apply softmax.

		Returns:
			a Tensor of shape [batch_size, out_channels, time_length]
		"""
		#[batch_size, channels, time_length] -> [batch_size, time_length, channels]
		batch_size = tf.shape(x)[0]
		time_length = tf.shape(x)[-1]

		if g is not None:
			if self.embed_speakers is not None:
				#[batch_size, 1] ==> [batch_size, 1, gin_channels]
				g = self.embed_speakers(tf.reshape(g, [batch_size, -1]))
				#[batch_size, gin_channels, 1]
				with tf.control_dependencies([tf.assert_equal(tf.rank(g), 3)]):
					g = tf.transpose(g, [0, 2, 1])

		#Expand global conditioning features to all time steps
		g_bct = _expand_global_features(batch_size, time_length, g, data_format='BCT')

		if c is not None and self.upsample_conv is not None:
			#[batch_size, 1, cin_channels, time_length]
			c = tf.expand_dims(c, axis=1)
			for transposed_conv in self.upsample_conv:
				c = transposed_conv(c)

			#[batch_size, cin_channels, time_length]
			c = tf.squeeze(c, [1])
			with tf.control_dependencies([tf.assert_equal(tf.shape(c)[-1], tf.shape(x)[-1])]):
				c = tf.identity(c, name='control_c_and_x_shape')

		#Feed data to network
		x = self.first_conv(x)
		skips = None
		for conv in self.conv_layers:
			x, h = conv(x, c, g_bct)
			if skips is None:
				skips = h
			else:
				skips = skips + h
		x = skips

		for conv in self.last_conv_layers:
			x = conv(x)

		return tf.nn.softmax(x, axis=1) if softmax else x