Python tensorflow.compat.v2.Tensor() Examples

The following are 30 code examples of tensorflow.compat.v2.Tensor(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module tensorflow.compat.v2 , or try the search function .
Example #1
Source File: date_tensor.py    From tf-quant-finance with Apache License 2.0 6 votes vote down vote up
def to_tensor(self):
    """Packs the dates into a single Tensor.

    The Tensor has shape `date_tensor.shape() + (3,)`, where the last dimension
    represents years, months and days, in this order.

    This can be convenient when the dates are the final result of a computation
    in the graph mode: a `tf.function` can return `date_tensor.to_tensor()`, or,
    if one uses `tf.compat.v1.Session`, they can call
    `session.run(date_tensor.to_tensor())`.

    Returns:
      A Tensor of shape `date_tensor.shape() + (3,)`.

    #### Example

    ```python
    dates = tff.datetime.dates_from_tuples([(2019, 1, 25), (2020, 3, 2)])
    dates.to_tensor()  # tf.Tensor with contents [[2019, 1, 25], [2020, 3, 2]].
    ```
    """
    return tf.stack((self.year(), self.month(), self.day()), axis=-1) 
Example #2
Source File: array_ops.py    From trax with Apache License 2.0 6 votes vote down vote up
def imag(a):
  """Returns imaginary parts of all elements in `a`.

  Uses `tf.imag`.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.

  Returns:
    An ndarray with the same shape as `a`.
  """
  a = asarray(a)
  # TODO(srbs): np.imag returns a scalar if a is a scalar, whereas we always
  # return an ndarray.
  return utils.tensor_to_ndarray(tf.math.imag(a.data)) 
Example #3
Source File: array_ops.py    From trax with Apache License 2.0 6 votes vote down vote up
def all(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
  """Whether all array elements or those along an axis evaluate to true.

  Casts the array to bool type if it is not already and uses `tf.reduce_all` to
  compute the result.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    axis: Optional. Could be an int or a tuple of integers. If not specified,
      the reduction is performed over all array indices.
    keepdims: If true, retains reduced dimensions with length 1.

  Returns:
    An ndarray. Note that unlike NumPy this does not return a scalar bool if
    `axis` is None.
  """
  a = asarray(a, dtype=bool)
  return utils.tensor_to_ndarray(
      tf.reduce_all(input_tensor=a.data, axis=axis, keepdims=keepdims)) 
Example #4
Source File: continuous_batched.py    From compression with Apache License 2.0 6 votes vote down vote up
def quantize(self, bottleneck):
    """Quantizes a floating-point tensor.

    To use this entropy model as an information bottleneck during training, pass
    a tensor through this function. The tensor is rounded to integer values
    shifted by an offset, which depends on `self.prior`. For instance, for a
    Gaussian distribution, the returned values are rounded to the location of
    the mode of the distribution plus or minus an integer.

    The gradient of this rounding operation is overridden with the identity
    (straight-through gradient estimator).

    Arguments:
      bottleneck: `tf.Tensor` containing the data to be quantized. The innermost
        dimensions must be broadcastable to `self.prior_shape`.

    Returns:
      A `tf.Tensor` containing the quantized values.
    """
    return self._quantize(bottleneck, self._quantization_offset) 
Example #5
Source File: array_ops.py    From trax with Apache License 2.0 6 votes vote down vote up
def diagflat(v, k=0):
  """Returns a 2-d array with flattened `v` as diagonal.

  Args:
    v: array_like of any rank. Gets flattened when setting as diagonal. Could be
      an ndarray, a Tensor or any object that can be converted to a Tensor using
      `tf.convert_to_tensor`.
    k: Position of the diagonal. Defaults to 0, the main diagonal. Positive
      values refer to diagonals shifted right, negative values refer to
      diagonals shifted left.

  Returns:
    2-d ndarray.
  """
  v = asarray(v)
  return diag(tf.reshape(v.data, [-1]), k) 
Example #6
Source File: array_ops.py    From trax with Apache License 2.0 6 votes vote down vote up
def real(val):
  """Returns real parts of all elements in `a`.

  Uses `tf.real`.

  Args:
    val: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.

  Returns:
    An ndarray with the same shape as `a`.
  """
  val = asarray(val)
  # TODO(srbs): np.real returns a scalar if val is a scalar, whereas we always
  # return an ndarray.
  return utils.tensor_to_ndarray(tf.math.real(val.data)) 
Example #7
Source File: array_ops.py    From trax with Apache License 2.0 6 votes vote down vote up
def ones_like(a, dtype=None):
  """Returns an array of ones with the shape and type of the input array.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can be
      converted to a Tensor using `tf.convert_to_tensor`.
    dtype: Optional, defaults to dtype of the input array. The type of the
      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
      `DType`.

  Returns:
    An ndarray.
  """
  if isinstance(a, arrays_lib.ndarray):
    a = a.data
  if dtype is None:
    dtype = utils.result_type(a)
  else:
    dtype = utils.result_type(dtype)
  return arrays_lib.tensor_to_ndarray(tf.ones_like(a, dtype)) 
Example #8
Source File: model_tf2.py    From machine-learning-for-programming-samples with MIT License 6 votes vote down vote up
def compute_logits(self, token_ids: tf.Tensor, training: bool) -> tf.Tensor:
        """
        Implements a language model, where each output is conditional on the current
        input and inputs processed so far.

        Args:
            token_ids: int32 tensor of shape [B, T], storing integer IDs of tokens.
            training: Flag indicating if we are currently training (used to toggle dropout)

        Returns:
            tf.float32 tensor of shape [B, T, V], storing the distribution over output symbols
            for each timestep for each batch element.
        """
        # TODO 5# 1) Embed tokens
        # TODO 5# 2) Run RNN on embedded tokens
        # TODO 5# 3) Project RNN outputs onto the vocabulary to obtain logits.
        return rnn_output_logits 
Example #9
Source File: array_ops.py    From trax with Apache License 2.0 6 votes vote down vote up
def squeeze(a, axis=None):
  """Removes single-element axes from the array.

  Args:
    a: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    axis: scalar or list/tuple of ints.

  TODO(srbs): tf.squeeze throws error when axis is a Tensor eager execution
  is enabled. So we cannot allow axis to be array_like here. Fix.

  Returns:
    An ndarray.
  """
  a = asarray(a)
  return utils.tensor_to_ndarray(tf.squeeze(a, axis)) 
Example #10
Source File: utils.py    From trax with Apache License 2.0 6 votes vote down vote up
def result_type(*arrays_and_dtypes):
  """Returns the type resulting from applying NumPy type promotion to arguments.

  Args:
    *arrays_and_dtypes: A list of array_like objects or dtypes.

  Returns:
    A numpy dtype.
  """
  def maybe_get_dtype(x):
    # Don't put np.ndarray in this list, because np.result_type looks at the
    # value (not just dtype) of np.ndarray to decide the result type.
    if isinstance(x, (arrays.ndarray, arrays.ShardedNdArray,
                      tf.Tensor, tf.IndexedSlices)):
      return _to_numpy_type(x.dtype)
    elif isinstance(x, tf.DType):
      return _to_numpy_type(x)
    return x
  arrays_and_dtypes = [maybe_get_dtype(x) for x in
                       tf.nest.flatten(arrays_and_dtypes)]
  if not arrays_and_dtypes:
    # If arrays_and_dtypes is an empty list, let numpy decide what the dtype is.
    arrays_and_dtypes = [np.asarray([])]
  return dtypes._result_type(*arrays_and_dtypes) 
Example #11
Source File: multi_objective_scalarizer.py    From agents with Apache License 2.0 6 votes vote down vote up
def __call__(self, multi_objectives: tf.Tensor) -> tf.Tensor:
    """Returns a single reward by scalarizing multiple objectives.

    Args:
      multi_objectives: A `Tensor` of shape [batch_size, number_of_objectives],
        where each column represents an objective.

    Returns: A `Tensor` of shape [batch_size] representing scalarized rewards.

    Raises:
      ValueError: if `multi_objectives.shape.rank != 2`.
      ValueError: if
        `multi_objectives.shape.dims[1] != self._num_of_objectives`.
    """
    if multi_objectives.shape.rank != 2:
      raise ValueError('The rank of the input should be 2, but is {}'.format(
          multi_objectives.shape.rank))
    if multi_objectives.shape.dims[1] != self._num_of_objectives:
      raise ValueError(
          'The number of input objectives should be {}, but is {}.'.format(
              self._num_of_objectives, multi_objectives.shape.dims[1]))
    return self.call(multi_objectives)

  # Subclasses must implement these methods. 
Example #12
Source File: test_util.py    From spectral-density with Apache License 2.0 6 votes vote down vote up
def hessian(function: Callable[[Parameters], tf.Tensor],
            parameters: Parameters) -> Parameters:
  """Computes the Hessian of a given function.

  Useful for testing, although scales very poorly.

  Args:
    function: A function for which we want to compute the Hessian.
    parameters: Parameters with respect to the Hessian should be computed.

  Returns:
    A tensor or list of tensors of same nested structure as `Parameters`,
      representing the Hessian.
  """
  with tf.GradientTape() as outer_tape:
    with tf.GradientTape() as inner_tape:
      value = function(parameters)
    grads = inner_tape.gradient(value, parameters)
    grads = tensor_list_util.tensor_list_to_vector(grads)
  return outer_tape.jacobian(grads, parameters) 
Example #13
Source File: test_util.py    From spectral-density with Apache License 2.0 6 votes vote down vote up
def hessian_as_matrix(function: Callable[[Parameters], tf.Tensor],
                      parameters: Parameters) -> tf.Tensor:
  """Computes the Hessian of a given function.

  Same as `hessian`, although return a matrix of size [w_dim, w_dim], where
  `w_dim` is the number of parameters, which makes it easier to work with.

  Args:
    function: A function for which we want to compute the Hessian.
    parameters: Parameters with respect to the Hessian should be computed.

  Returns:
    A tensor of size [w_dim, w_dim] representing the Hessian.
  """
  hessian_as_tensor_list = hessian(function, parameters)
  hessian_as_tensor_list = [
      tf.reshape(e, [e.shape[0], -1]) for e in hessian_as_tensor_list]
  return tf.concat(hessian_as_tensor_list, axis=1) 
Example #14
Source File: feature.py    From ranking with Apache License 2.0 6 votes vote down vote up
def __init__(self,
               example_feature_columns,
               size_feature_name,
               name='generate_mask_layer',
               **kwargs):
    """Constructs a mask generator layer.

    Args:
      example_feature_columns: (dict) example feature names to columns.
      size_feature_name: (str) Name of feature for example list sizes. If not
        None, this feature name corresponds to a `tf.int32` Tensor of size
        [batch_size] corresponding to sizes of example lists. If `None`, all
        examples are treated as valid.
      name: (str) name of the layer.
      **kwargs: keyword arguments.
    """
    super(GenerateMask, self).__init__(name=name, **kwargs)
    self._example_feature_columns = example_feature_columns
    self._size_feature_name = size_feature_name 
Example #15
Source File: feature.py    From ranking with Apache License 2.0 6 votes vote down vote up
def call(self, inputs):
    """Generates mask (whether example is valid) from features.

    Args:
      inputs: (dict) Features with a mix of context (2D) and example features
        (3D).

    Returns:
      mask: (tf.Tensor) Mask is a tensor of shape [batch_size, list_size], which
        is True for a valid example and False for invalid one.
    """
    example_feature = inputs[next(six.iterkeys(self._example_feature_columns))]
    list_size = tf.shape(example_feature)[1]
    sizes = inputs[self._size_feature_name]
    mask = tf.sequence_mask(sizes, maxlen=list_size)
    return mask 
Example #16
Source File: network.py    From ranking with Apache License 2.0 6 votes vote down vote up
def transform(self, features=None, training=None, mask=None):
    """Transforms the features into dense context features and example features.

    The user can overwrite this function for custom transformations.
    Mask is provided as an argument so that inherited models can have access
    to it for custom feature transformations, without modifying
    `call` explicitly.

    Args:
      features: (dict) with a mix of context (2D) and example features (3D).
      training: (bool) whether in train or inference mode.
      mask: (tf.Tensor) Mask is a tensor of shape [batch_size, list_size], which
        is True for a valid example and False for invalid one.

    Returns:
      context_features: (dict) context feature names to dense 2D tensors of
        shape [batch_size, feature_dims].
      example_features: (dict) example feature names to dense 3D tensors of
        shape [batch_size, list_size, feature_dims].
    """
    del mask
    context_features, example_features = self._listwise_dense_layer(
        inputs=features, training=training)
    return context_features, example_features 
Example #17
Source File: network.py    From ranking with Apache License 2.0 6 votes vote down vote up
def call(self, inputs=None, training=None, mask=None):
    """Defines the forward pass for ranking model.

    Args:
      inputs: (dict) with a mix of context (2D) and example features (3D).
      training: (bool) whether in train or inference mode.
      mask: (tf.Tensor) Mask is a tensor of shape [batch_size, list_size], which
        is True for a valid example and False for invalid one.

    Returns:
      (tf.Tensor) A score tensor of shape [batch_size, list_size].
    """
    context_features, example_features = self.transform(
        features=inputs, training=training, mask=mask)
    logits = self.compute_logits(
        context_features=context_features,
        example_features=example_features,
        training=training,
        mask=mask)
    return logits 
Example #18
Source File: custom_loops.py    From tf-quant-finance with Apache License 2.0 6 votes vote down vote up
def _block_matmul(m1, m2):
  """Multiplies block matrices represented as nested lists."""
  # Calls itself recursively to multiply blocks, until reaches the level of
  # tf.Tensors.
  if isinstance(m1, tf.Tensor):
    assert isinstance(m2, tf.Tensor)
    return tf.matmul(m1, m2)
  assert _is_nested_list(m1) and _is_nested_list(m2)

  i_max = len(m1)
  k_max = len(m2)
  j_max = 0 if k_max == 0 else len(m2[0])
  if i_max > 0:
    assert len(m1[0]) == k_max

  def row_by_column(i, j):
    return _block_add(*[_block_matmul(m1[i][k], m2[k][j])
                        for k in range(k_max)])
  return [[row_by_column(i, j) for j in range(j_max)] for i in range(i_max)] 
Example #19
Source File: custom_loops.py    From tf-quant-finance with Apache License 2.0 6 votes vote down vote up
def _block_add(*ms):
  """Adds block matrices represented as nested lists."""
  # Calls itself recursively to add blocks, until reaches the level of
  # tf.Tensors.
  if len(ms) == 1:
    return ms[0]
  if isinstance(ms[0], tf.Tensor):
    assert all(isinstance(m, tf.Tensor) for m in ms[1:])
    return tf.math.add_n(ms)
  assert all(_is_nested_list(m) for m in ms)
  for i in range(1, len(ms)):
    tf.nest.assert_same_structure(ms[0], ms[i])

  i_max = len(ms[0])
  j_max = 0 if i_max == 0 else len(ms[0][0])
  return [[_block_add(*[ms[k][i][j] for k in range(len(ms))])
           for j in range(j_max)]
          for i in range(i_max)] 
Example #20
Source File: __init__.py    From language with Apache License 2.0 6 votes vote down vote up
def matmul_any_tensor_dense_tensor(a,
                                   b,
                                   a_is_sparse = True,
                                   transpose_a = False):
  """Like tf.matmul except that first argument might be a SparseTensor.

  Args:
    a: SparseTensor or Tensor
    b: Tensor
    a_is_sparse: true if a is a SparseTensor
    transpose_a: transpose a before performing matmal operation.

  Returns:
    TF expression for a.dot(b) or a.transpose().dot(b)

  Raises
    ValueError: If a or b are of the wrong type.
  """
  if a_is_sparse:
    _check_type('a', a, tf.SparseTensor)
    return tf.sparse.sparse_dense_matmul(
        b, a, adjoint_a=False, adjoint_b=not transpose_a)
  else:
    return tf.transpose(
        a=tf.matmul(a, tf.transpose(a=b), transpose_a=transpose_a)) 
Example #21
Source File: inner_reshape.py    From agents with Apache License 2.0 6 votes vote down vote up
def _reshape_inner_dims(
    tensor: tf.Tensor,
    shape: tf.TensorShape,
    new_shape: tf.TensorShape) -> tf.Tensor:
  """Reshapes tensor to: shape(tensor)[:-len(shape)] + new_shape."""
  tensor_shape = tf.shape(tensor)
  ndims = shape.rank
  tensor.shape[-ndims:].assert_is_compatible_with(shape)
  new_shape_inner_tensor = tf.cast(
      [-1 if d is None else d for d in new_shape.as_list()], tf.int64)
  new_shape_outer_tensor = tf.cast(
      tensor_shape[:-ndims], tf.int64)
  full_new_shape = tf.concat(
      (new_shape_outer_tensor, new_shape_inner_tensor), axis=0)
  new_tensor = tf.reshape(tensor, full_new_shape)
  new_tensor.set_shape(tensor.shape[:-ndims] + new_shape)
  return new_tensor 
Example #22
Source File: continuous_indexed.py    From compression with Apache License 2.0 5 votes vote down vote up
def bits(self, bottleneck, indexes, training=True):
    """Estimates the number of bits needed to compress a tensor.

    Arguments:
      bottleneck: `tf.Tensor` containing the data to be compressed.
      indexes: `tf.Tensor` specifying the scalar distribution for each element
        in `bottleneck`. See class docstring for examples.
      training: Boolean. If `False`, computes the Shannon information of
        `bottleneck` under the distribution computed by `self.prior_fn`,
        which is a non-differentiable, tight *lower* bound on the number of bits
        needed to compress `bottleneck` using `compress()`. If `True`, returns a
        somewhat looser, but differentiable *upper* bound on this quantity.

    Returns:
      A `tf.Tensor` having the same shape as `bottleneck` without the
      `self.coding_rank` innermost dimensions, containing the number of bits.
    """
    indexes = self._normalize_indexes(indexes)
    prior = self._make_prior(indexes)
    if training:
      quantized = bottleneck + tf.random.uniform(
          tf.shape(bottleneck), minval=-.5, maxval=.5, dtype=bottleneck.dtype)
    else:
      offset = helpers.quantization_offset(prior)
      quantized = self._quantize(bottleneck, offset)
    probs = prior.prob(quantized)
    probs = math_ops.lower_bound(probs, self.likelihood_bound)
    axes = tuple(range(-self.coding_rank, 0))
    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / -tf.math.log(2.)
    return bits 
Example #23
Source File: date_tensor.py    From tf-quant-finance with Apache License 2.0 5 votes vote down vote up
def is_end_of_month(self):
    """Returns a bool Tensor indicating whether dates are at ends of months."""
    return tf.math.equal(self._days,
                         _num_days_in_month(self._months, self._years)) 
Example #24
Source File: date_tensor.py    From tf-quant-finance with Apache License 2.0 5 votes vote down vote up
def __eq__(self, other):
    """Compares two DateTensors by "==", returning a Tensor of bools."""
    # Note that tf doesn't override "==" and  "!=", unlike numpy.
    return tf.math.equal(self._ordinals, other.ordinal()) 
Example #25
Source File: date_tensor.py    From tf-quant-finance with Apache License 2.0 5 votes vote down vote up
def day_of_year(self):
    """Calculates the number of days since the beginning of the year.

    Returns:
      Tensor of int32 type with elements in range [1, 366]. January 1st yields
      "1".

    #### Example

    ```python
    dt = tff.datetime.dates_from_tuples([(2019, 1, 25), (2020, 3, 2)])
    dt.day_of_year()  # [25, 62]
    ```
    """
    if self._day_of_year is None:
      cumul_days_in_month_nonleap = tf.math.cumsum(
          _DAYS_IN_MONTHS_NON_LEAP, exclusive=True)
      cumul_days_in_month_leap = tf.math.cumsum(
          _DAYS_IN_MONTHS_LEAP, exclusive=True)
      days_before_month_non_leap = tf.gather(cumul_days_in_month_nonleap,
                                             self.month() - 1)
      days_before_month_leap = tf.gather(cumul_days_in_month_leap,
                                         self.month() - 1)
      days_before_month = tf.where(
          date_utils.is_leap_year(self.year()), days_before_month_leap,
          days_before_month_non_leap)
      self._day_of_year = days_before_month + self.day()
    return self._day_of_year 
Example #26
Source File: uniform_noise.py    From compression with Apache License 2.0 5 votes vote down vote up
def _logsum_expbig_minus_expsmall(big, small):
  """Numerically stable evaluation of `log(exp(big) - exp(small))`.

  This assumes `small <= big` and arguments that can be broadcast against each
  other.

  Arguments:
    big: Floating-point `tf.Tensor`.
    small: Floating-point `tf.Tensor`.

  Returns:
    `tf.Tensor` containing the result.
  """
  with tf.name_scope("logsum_expbig_minus_expsmall"):
    return tf.math.log1p(-tf.exp(small - big)) + big 
Example #27
Source File: helpers.py    From compression with Apache License 2.0 5 votes vote down vote up
def upper_tail(distribution, tail_mass):
  """Approximates upper tail quantile for range coding.

  For range coding of random variables, the distribution tails need special
  handling, because range coding can only handle alphabets with a finite
  number of symbols. This method returns a cut-off location for the upper
  tail, such that approximately `tail_mass` probability mass is contained in
  the tails (together). The tails are then handled by using the 'overflow'
  functionality of the range coder implementation (using a Golomb-like
  universal code).

  Arguments:
    distribution: A `tfp.distributions.Distribution` object.
    tail_mass: Float between 0 and 1. Desired probability mass for the tails.

  Returns:
    A `tf.Tensor` broadcastable to shape `self.batch_shape` containing the
    approximate upper tail quantiles for each scalar distribution.
  """
  try:
    tail = distribution._upper_tail(tail_mass)  # pylint:disable=protected-access
  except (AttributeError, NotImplementedError):
    try:
      tail = distribution.quantile(1 - tail_mass / 2)
    except NotImplementedError:
      try:
        tail = estimate_tails(
            distribution.log_survival_function, tf.math.log(tail_mass / 2),
            distribution.batch_shape_tensor(), distribution.dtype)
      except NotImplementedError:
        raise NotImplementedError(
            "`distribution` must implement `_upper_tail()`, `quantile()`, or "
            "`log_survival_function()` so that upper tail can be located.")
  return tf.stop_gradient(tail) 
Example #28
Source File: holiday_calendar_factory.py    From tf-quant-finance with Apache License 2.0 5 votes vote down vote up
def _tensor_is_not_empty(t):
  """Returns whether t is definitely not empty."""
  # False means either empty or unknown.
  if t is None:
    return False
  if isinstance(t, np.ndarray):
    return t.size > 0
  if isinstance(t, tf.Tensor):
    num_elem = t.shape.num_elements
    return num_elem is not None and num_elem > 0  # None means shape is unknown.
  return bool(t) 
Example #29
Source File: date_tensor.py    From tf-quant-finance with Apache License 2.0 5 votes vote down vote up
def __lt__(self, other):
    """Compares two DateTensors by "<", returning a Tensor of bools."""

    return self._ordinals < other.ordinal() 
Example #30
Source File: continuous_batched.py    From compression with Apache License 2.0 5 votes vote down vote up
def bits(self, bottleneck, training=True):
    """Estimates the number of bits needed to compress a tensor.

    Arguments:
      bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
        least `self.coding_rank` dimensions, and the innermost dimensions must
        be broadcastable to `self.prior_shape`.
      training: Boolean. If `False`, computes the Shannon information of
        `bottleneck` under the distribution `self.prior`, which is a
        non-differentiable, tight *lower* bound on the number of bits needed to
        compress `bottleneck` using `compress()`. If `True`, returns a somewhat
        looser, but differentiable *upper* bound on this quantity.

    Returns:
      A `tf.Tensor` having the same shape as `bottleneck` without the
      `self.coding_rank` innermost dimensions, containing the number of bits.
    """
    if training:
      quantized = bottleneck + tf.random.uniform(
          tf.shape(bottleneck), minval=-.5, maxval=.5, dtype=bottleneck.dtype)
    else:
      quantized = self.quantize(bottleneck)
    probs = self.prior.prob(quantized)
    probs = math_ops.lower_bound(probs, self.likelihood_bound)
    axes = tuple(range(-self.coding_rank, 0))
    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / -tf.math.log(2.)
    return bits