Skip to content

Commit 73fd15e

Browse files
feat: double precision for ML backends (#701)
* Add option to use double precision in ML backends
1 parent 62becc2 commit 73fd15e

File tree

6 files changed

+37
-29
lines changed

6 files changed

+37
-29
lines changed

src/pyhf/cli/infer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,11 +66,11 @@ def cls(
6666

6767
# set the backend if not NumPy
6868
if backend in ['pytorch', 'torch']:
69-
set_backend(tensor.pytorch_backend())
69+
set_backend(tensor.pytorch_backend(float='float64'))
7070
elif backend in ['tensorflow', 'tf']:
7171
from tensorflow.compat.v1 import Session
7272

73-
set_backend(tensor.tensorflow_backend(session=Session()))
73+
set_backend(tensor.tensorflow_backend(session=Session(), float='float64'))
7474
tensorlib, _ = get_backend()
7575

7676
optconf = {k: v for item in optconf for k, v in item.items()}

src/pyhf/optimize/opt_tflow.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,10 @@ def setup_minimize(
3838
variable_bounds = [par_bounds[i] for i in variable_idx]
3939

4040
data_placeholder = tf.placeholder(
41-
tf.float32, (pdf.config.nmaindata + pdf.config.nauxdata,)
41+
tensorlib.dtypemap['float'], (pdf.config.nmaindata + pdf.config.nauxdata,)
4242
)
4343
variable_pars_placeholder = tf.placeholder(
44-
tf.float32, (pdf.config.npars - len(fixed_vals),)
44+
tensorlib.dtypemap['float'], (pdf.config.npars - len(fixed_vals),)
4545
)
4646

4747
tv = _TensorViewer([fixed_idx, variable_idx])

src/pyhf/tensor/numpy_backend.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
"""NumPy Tensor Library Module."""
12
import numpy as np
23
import logging
34
from scipy.special import gammaln

src/pyhf/tensor/pytorch_backend.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
"""PyTorch Tensor Library Module."""
12
import torch
23
import torch.autograd
34
import logging
@@ -10,6 +11,11 @@ class pytorch_backend(object):
1011

1112
def __init__(self, **kwargs):
1213
self.name = 'pytorch'
14+
self.dtypemap = {
15+
'float': getattr(torch, kwargs.get('float', 'float32')),
16+
'int': getattr(torch, kwargs.get('float', 'int32')),
17+
'bool': torch.bool,
18+
}
1319

1420
def clip(self, tensor_in, min_value, max_value):
1521
"""
@@ -100,9 +106,8 @@ def astensor(self, tensor_in, dtype='float'):
100106
Returns:
101107
torch.Tensor: A multi-dimensional matrix containing elements of a single data type.
102108
"""
103-
dtypemap = {'float': torch.float, 'int': torch.int, 'bool': torch.bool}
104109
try:
105-
dtype = dtypemap[dtype]
110+
dtype = self.dtypemap[dtype]
106111
except KeyError:
107112
log.error('Invalid dtype: dtype must be float, int, or bool.')
108113
raise
@@ -141,10 +146,10 @@ def abs(self, tensor):
141146
return torch.abs(tensor)
142147

143148
def ones(self, shape):
144-
return torch.Tensor(torch.ones(shape))
149+
return torch.ones(shape, dtype=self.dtypemap['float'])
145150

146151
def zeros(self, shape):
147-
return torch.Tensor(torch.zeros(shape))
152+
return torch.zeros(shape, dtype=self.dtypemap['float'])
148153

149154
def power(self, tensor_in_1, tensor_in_2):
150155
return torch.pow(tensor_in_1, tensor_in_2)

src/pyhf/tensor/tensorflow_backend.py

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
"""Tensorflow Tensor Library Module."""
12
import logging
23
import tensorflow as tf
34
import tensorflow_probability as tfp
@@ -11,13 +12,17 @@ class tensorflow_backend(object):
1112
def __init__(self, **kwargs):
1213
self.session = kwargs.get('session')
1314
self.name = 'tensorflow'
15+
self.dtypemap = {
16+
'float': getattr(tf, kwargs.get('float', 'float32')),
17+
'int': getattr(tf, kwargs.get('int', 'int32')),
18+
'bool': tf.bool,
19+
}
1420

1521
def clip(self, tensor_in, min_value, max_value):
1622
"""
1723
Clips (limits) the tensor values to be within a specified min and max.
1824
1925
Example:
20-
2126
>>> import pyhf
2227
>>> import tensorflow as tf
2328
>>> sess = tf.compat.v1.Session()
@@ -36,6 +41,7 @@ def clip(self, tensor_in, min_value, max_value):
3641
3742
Returns:
3843
TensorFlow Tensor: A clipped `tensor`
44+
3945
"""
4046
if min_value is None:
4147
min_value = tf.reduce_min(tensor_in)
@@ -48,7 +54,6 @@ def tile(self, tensor_in, repeats):
4854
Repeat tensor data along a specific dimension
4955
5056
Example:
51-
5257
>>> import pyhf
5358
>>> import tensorflow as tf
5459
>>> sess = tf.compat.v1.Session()
@@ -67,6 +72,7 @@ def tile(self, tensor_in, repeats):
6772
6873
Returns:
6974
TensorFlow Tensor: The tensor with repeated axes
75+
7076
"""
7177
return tf.tile(tensor_in, repeats)
7278

@@ -75,7 +81,6 @@ def conditional(self, predicate, true_callable, false_callable):
7581
Runs a callable conditional on the boolean value of the evaulation of a predicate
7682
7783
Example:
78-
7984
>>> import pyhf
8085
>>> import tensorflow as tf
8186
>>> sess = tf.compat.v1.Session()
@@ -97,6 +102,7 @@ def conditional(self, predicate, true_callable, false_callable):
97102
98103
Returns:
99104
TensorFlow Tensor: The output of the callable that was evaluated
105+
100106
"""
101107
return tf.cond(predicate, true_callable, false_callable)
102108

@@ -157,10 +163,10 @@ def astensor(self, tensor_in, dtype='float'):
157163
158164
Returns:
159165
`tf.Tensor`: A symbolic handle to one of the outputs of a `tf.Operation`.
166+
160167
"""
161-
dtypemap = {'float': tf.float32, 'int': tf.int32, 'bool': tf.bool}
162168
try:
163-
dtype = dtypemap[dtype]
169+
dtype = self.dtypemap[dtype]
164170
except KeyError:
165171
log.error('Invalid dtype: dtype must be float, int, or bool.')
166172
raise
@@ -198,10 +204,10 @@ def abs(self, tensor):
198204
return tf.abs(tensor)
199205

200206
def ones(self, shape):
201-
return tf.ones(shape)
207+
return tf.ones(shape, dtype=self.dtypemap['float'])
202208

203209
def zeros(self, shape):
204-
return tf.zeros(shape)
210+
return tf.zeros(shape, dtype=self.dtypemap['float'])
205211

206212
def power(self, tensor_in_1, tensor_in_2):
207213
return tf.pow(tensor_in_1, tensor_in_2)
@@ -249,7 +255,6 @@ def simple_broadcast(self, *args):
249255
Broadcast a sequence of 1 dimensional arrays.
250256
251257
Example:
252-
253258
>>> import pyhf
254259
>>> import tensorflow as tf
255260
>>> sess = tf.compat.v1.Session()
@@ -266,6 +271,7 @@ def simple_broadcast(self, *args):
266271
267272
Returns:
268273
list of Tensors: The sequence broadcast together.
274+
269275
"""
270276
max_dim = max(map(lambda arg: arg.shape[0], args))
271277
try:
@@ -308,7 +314,6 @@ def poisson_logpdf(self, n, lam):
308314
at :code:`n` given the parameter :code:`lam`.
309315
310316
Example:
311-
312317
>>> import pyhf
313318
>>> import tensorflow as tf
314319
>>> sess = tf.compat.v1.Session()
@@ -343,7 +348,6 @@ def poisson(self, n, lam):
343348
at :code:`n` given the parameter :code:`lam`.
344349
345350
Example:
346-
347351
>>> import pyhf
348352
>>> import tensorflow as tf
349353
>>> sess = tf.compat.v1.Session()
@@ -378,7 +382,6 @@ def normal_logpdf(self, x, mu, sigma):
378382
of :code:`sigma`.
379383
380384
Example:
381-
382385
>>> import pyhf
383386
>>> import tensorflow as tf
384387
>>> sess = tf.compat.v1.Session()
@@ -414,7 +417,6 @@ def normal(self, x, mu, sigma):
414417
of :code:`sigma`.
415418
416419
Example:
417-
418420
>>> import pyhf
419421
>>> import tensorflow as tf
420422
>>> sess = tf.compat.v1.Session()
@@ -443,12 +445,11 @@ def normal(self, x, mu, sigma):
443445
normal = tfp.distributions.Normal(mu, sigma)
444446
return normal.prob(x)
445447

446-
def normal_cdf(self, x, mu=0, sigma=1):
448+
def normal_cdf(self, x, mu=0.0, sigma=1):
447449
"""
448-
The cumulative distribution function for the Normal distribution
450+
Compute the value of cumulative distribution function for the Normal distribution at x.
449451
450452
Example:
451-
452453
>>> import pyhf
453454
>>> import tensorflow as tf
454455
>>> sess = tf.compat.v1.Session()
@@ -472,15 +473,16 @@ def normal_cdf(self, x, mu=0, sigma=1):
472473
Returns:
473474
TensorFlow Tensor: The CDF
474475
"""
475-
normal = tfp.distributions.Normal(mu, sigma)
476+
normal = tfp.distributions.Normal(
477+
self.astensor(mu, dtype='float')[0], self.astensor(sigma, dtype='float')[0],
478+
)
476479
return normal.cdf(x)
477480

478481
def poisson_dist(self, rate):
479482
r"""
480-
The Poisson distribution with rate parameter :code:`rate`.
483+
Construct a Poisson distribution with rate parameter :code:`rate`.
481484
482485
Example:
483-
484486
>>> import pyhf
485487
>>> import tensorflow as tf
486488
>>> sess = tf.compat.v1.Session()
@@ -505,10 +507,9 @@ def poisson_dist(self, rate):
505507

506508
def normal_dist(self, mu, sigma):
507509
r"""
508-
The Normal distribution with mean :code:`mu` and standard deviation :code:`sigma`.
510+
Construct a Normal distribution with mean :code:`mu` and standard deviation :code:`sigma`.
509511
510512
Example:
511-
512513
>>> import pyhf
513514
>>> import tensorflow as tf
514515
>>> sess = tf.compat.v1.Session()

tests/conftest.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,13 +44,14 @@ def reset_backend():
4444
params=[
4545
(pyhf.tensor.numpy_backend(), None),
4646
(pyhf.tensor.pytorch_backend(), None),
47+
(pyhf.tensor.pytorch_backend(float='float64', int='int64'), None),
4748
(pyhf.tensor.tensorflow_backend(session=tf.compat.v1.Session()), None),
4849
(
4950
pyhf.tensor.numpy_backend(poisson_from_normal=True),
5051
pyhf.optimize.minuit_optimizer(),
5152
),
5253
],
53-
ids=['numpy', 'pytorch', 'tensorflow', 'numpy_minuit'],
54+
ids=['numpy', 'pytorch', 'pytorch64', 'tensorflow', 'numpy_minuit'],
5455
)
5556
def backend(request):
5657
# a better way to get the id? all the backends we have so far for testing

0 commit comments

Comments
 (0)