tensorflow / tensorflow/probability

Allow customizable dtype in discrete distributions

Open
#794 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Jupyter Notebook
Stars
4.4k
Forks
1.1k
No merged PRs in 30d

Description

I ran into something that I'm not sure if it was a deliberate design choice or if it's a small bug. I found that some discrete distributions cannot handle integer typed tensors in their log_prob method. The list of distributions is:

  • Poisson
  • NegativeBinomial
  • Geometric

Minimum working example that shows the problem

import tensorflow as tf
from tensorflow_probability import distributions as tfd

v_int = tf.constant(0, dtype="int32")
v_float = tf.constant(0, dtype="float32")

# The following work without raising an error
tfd.Poisson(rate=1.).log_prob(v_float)
tfd.Geometric(probs=0.5).log_prob(v_float)
tfd.NegativeBinomial(total_count=1, probs=0.5).log_prob(v_float)

# Any of the following fails raising different tracebacks
tfd.Poisson(rate=1.).log_prob(v_int)
tfd.Geometric(probs=0.5).log_prob(v_int)
tfd.NegativeBinomial(total_count=1, probs=0.5).log_prob(v_int)
`tfd.Poisson(rate=1.).log_prob(v_int)` Traceback --------------------------------------------------------------------------- TypeError Traceback (most recent call last) in ----> 1 tfd.Poisson(rate=1.).log_prob(v_int)

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in log_prob(self, value, name, **kwargs)
872 values of type self.dtype.
873 """
--> 874 return self._call_log_prob(value, name, **kwargs)
875
876 def _call_prob(self, value, name, **kwargs):

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in _call_log_prob(self, value, name, **kwargs)
854 value, name='value', dtype_hint=self.dtype)
855 if hasattr(self, '_log_prob'):
--> 856 return self._log_prob(value, **kwargs)
857 if hasattr(self, '_prob'):
858 return tf.math.log(self._prob(value, **kwargs))

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/poisson.py in _log_prob(self, x)
151 def _log_prob(self, x):
152 log_rate = self._log_rate_parameter_no_checks()
--> 153 log_probs = (self._log_unnormalized_prob(x, log_rate) -
154 self._log_normalization(log_rate))
155 if not self.interpolate_nondiscrete:

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/poisson.py in _log_unnormalized_prob(self, x, log_rate)
179 # The log-probability at negative points is always -inf.
180 # Catch such x's and set the output value accordingly.
--> 181 safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x), 0.)
182 y = safe_x * log_rate - tf.math.lgamma(1. + safe_x)
183 return tf.where(

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py in maximum(x, y, name)
5719 _result = pywrap_tfe.TFE_Py_FastPathExecute(
5720 _ctx._context_handle, tld.device_name, "Maximum", name,
-> 5721 tld.op_callbacks, x, y)
5722 return _result
5723 except _core._FallbackException:

TypeError: Cannot convert 0.0 to EagerTensor of dtype int32

`tfd.Geometric(probs=0.5).log_prob(v_int)` Traceback --------------------------------------------------------------------------- NotFoundError Traceback (most recent call last) in ----> 1 tfd.Geometric(probs=0.5).log_prob(v_int)

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in log_prob(self, value, name, **kwargs)
872 values of type self.dtype.
873 """
--> 874 return self._call_log_prob(value, name, **kwargs)
875
876 def _call_prob(self, value, name, **kwargs):

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in _call_log_prob(self, value, name, **kwargs)
854 value, name='value', dtype_hint=self.dtype)
855 if hasattr(self, '_log_prob'):
--> 856 return self._log_prob(value, **kwargs)
857 if hasattr(self, '_prob'):
858 return tf.math.log(self._prob(value, **kwargs))

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/geometric.py in _log_prob(self, x)
166 if not self.validate_args:
167 # For consistency with cdf, we take the floor.
--> 168 x = tf.floor(x)
169 safe_domain = tf.where(
170 tf.equal(x, 0.), tf.zeros_like(probs), probs)

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py in floor(x, name)
3720 raise
3721 except _core._NotOkStatusException as e:
-> 3722 _ops.raise_from_not_ok_status(e, name)
3723 # Add nodes to the TensorFlow graph.
3724 try:

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in raise_from_not_ok_status(e, name)
6623 message = e.message + (" name: " + name if name is not None else "")
6624 # pylint: disable=protected-access
-> 6625 six.raise_from(core._status_to_exception(e.code, message), None)
6626 # pylint: enable=protected-access
6627

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

NotFoundError: Could not find valid device for node.
Node:{{node Floor}}
All kernels registered for op Floor :
device='XLA_GPU'; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
device='XLA_CPU'; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
device='XLA_CPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
device='XLA_GPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
device='GPU'; T in [DT_DOUBLE]
device='GPU'; T in [DT_HALF]
device='GPU'; T in [DT_FLOAT]
device='CPU'; T in [DT_DOUBLE]
device='CPU'; T in [DT_HALF]
device='CPU'; T in [DT_FLOAT]
[Op:Floor]

`tfd.NegativeBinomial(total_count=1, probs=0.5).log_prob(v_int)` Traceback --------------------------------------------------------------------------- InvalidArgumentError Traceback (most recent call last) in ----> 1 tfd.NegativeBinomial(total_count=1, probs=0.5).log_prob(v_int)

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in log_prob(self, value, name, **kwargs)
872 values of type self.dtype.
873 """
--> 874 return self._call_log_prob(value, name, **kwargs)
875
876 def _call_prob(self, value, name, **kwargs):

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in _call_log_prob(self, value, name, **kwargs)
854 value, name='value', dtype_hint=self.dtype)
855 if hasattr(self, '_log_prob'):
--> 856 return self._log_prob(value, **kwargs)
857 if hasattr(self, '_prob'):
858 return tf.math.log(self._prob(value, **kwargs))

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_probability/python/distributions/negative_binomial.py in _log_prob(self, x)
185 logits = self._logits_parameter_no_checks()
186 log_unnormalized_prob = (total_count * tf.math.log_sigmoid(-logits) +
--> 187 x * tf.math.log_sigmoid(logits))
188 log_normalization = (-tf.math.lgamma(total_count + x) +
189 tf.math.lgamma(1. + x) +

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py in binary_op_wrapper(x, y)
943 with ops.name_scope(None, op_name, [x, y]) as name:
944 if isinstance(x, ops.Tensor) and isinstance(y, ops.Tensor):
--> 945 return func(x, y, name=name)
946 elif not isinstance(y, sparse_tensor.SparseTensor):
947 try:

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py in _mul_dispatch(x, y, name)
1242 is_tensor_y = isinstance(y, ops.Tensor)
1243 if is_tensor_y:
-> 1244 return gen_math_ops.mul(x, y, name=name)
1245 else:
1246 assert isinstance(y, sparse_tensor.SparseTensor) # Case: Dense * Sparse.

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py in mul(x, y, name)
6087 pass # Add nodes to the TensorFlow graph.
6088 except _core._NotOkStatusException as e:
-> 6089 _ops.raise_from_not_ok_status(e, name)
6090 # Add nodes to the TensorFlow graph.
6091 _, _, _op, _outputs = _op_def_library._apply_op_helper(

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in raise_from_not_ok_status(e, name)
6623 message = e.message + (" name: " + name if name is not None else "")
6624 # pylint: disable=protected-access
-> 6625 six.raise_from(core._status_to_exception(e.code, message), None)
6626 # pylint: enable=protected-access
6627

~/anaconda3/envs/pymc4/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a int32 tensor but is a float tensor [Op:Mul]

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the Poisson, Geometric, and NegativeBinomial distribution implementations named in the traceback, then reproduce the minimum working example using an int32 value. Done means each distribution's log_prob accepts integer-typed tensors without the shown errors while preserving the existing float behavior; add or update tests for these three cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.