1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
"""Required functions for optimized contractions of numpy arrays using tensorflow."""
from opt_einsum.helpers import has_array_interface
from opt_einsum.sharing import to_backend_cache_wrap
__all__ = ["to_tensorflow", "build_expression", "evaluate_constants"]
_CACHED_TF_DEVICE = None
def _get_tensorflow_and_device():
global _CACHED_TF_DEVICE
if _CACHED_TF_DEVICE is None:
import tensorflow as tf # type: ignore
try:
eager = tf.executing_eagerly()
except AttributeError:
try:
eager = tf.contrib.eager.in_eager_mode()
except AttributeError:
eager = False
device = tf.test.gpu_device_name()
if not device:
device = "cpu"
_CACHED_TF_DEVICE = tf, device, eager
return _CACHED_TF_DEVICE
@to_backend_cache_wrap(constants=True)
def to_tensorflow(array, constant=False):
"""Convert a numpy array to a ``tensorflow.placeholder`` instance."""
tf, device, eager = _get_tensorflow_and_device()
if eager:
if has_array_interface(array):
with tf.device(device):
return tf.convert_to_tensor(array)
return array
if has_array_interface(array):
if constant:
return tf.convert_to_tensor(array)
return tf.placeholder(array.dtype, array.shape)
return array
# Standard graph mode
def build_expression_graph(arrays, expr):
"""Build a tensorflow function based on ``arrays`` and ``expr``."""
tf, _, _ = _get_tensorflow_and_device()
placeholders = [to_tensorflow(array) for array in arrays]
graph = expr._contract(placeholders, backend="tensorflow")
def tensorflow_contract(*arrays):
session = tf.get_default_session()
# only want to feed placeholders - constant tensors already have values
feed_dict = {p: a for p, a in zip(placeholders, arrays) if p.op.type == "Placeholder"}
return session.run(graph, feed_dict=feed_dict)
return tensorflow_contract
def evaluate_constants_graph(const_arrays, expr):
"""Convert constant arguments to tensorflow constants, and perform any
possible constant contractions. Requires evaluating a tensorflow graph.
"""
tf, _, _ = _get_tensorflow_and_device()
# compute the partial graph of new inputs
const_arrays = [to_tensorflow(x, constant=True) for x in const_arrays]
new_ops, new_contraction_list = expr(*const_arrays, backend="tensorflow", evaluate_constants=True)
# evaluate the new inputs and convert back to tensorflow, maintaining None as non-consts
session = tf.get_default_session()
new_consts = iter(session.run([x for x in new_ops if x is not None]))
new_ops = [None if x is None else to_tensorflow(next(new_consts), constant=True) for x in new_ops]
return new_ops, new_contraction_list
# Eager execution mode
def build_expression_eager(_, expr):
"""Build a eager tensorflow function based on ``arrays`` and ``expr``."""
def tensorflow_eager_contract(*arrays):
return expr._contract([to_tensorflow(x) for x in arrays], backend="tensorflow").numpy()
return tensorflow_eager_contract
def evaluate_constants_eager(const_arrays, expr):
"""Convert constant arguments to tensorflow_eager arrays, and perform any
possible constant contractions.
"""
return expr(*[to_tensorflow(x) for x in const_arrays], backend="tensorflow", evaluate_constants=True)
# Dispatch to eager or graph mode
def build_expression(arrays, expr):
_, _, eager = _get_tensorflow_and_device()
fn = build_expression_eager if eager else build_expression_graph
return fn(arrays, expr)
def evaluate_constants(const_arrays, expr):
_, _, eager = _get_tensorflow_and_device()
fn = evaluate_constants_eager if eager else evaluate_constants_graph
return fn(const_arrays, expr)
|
