A configuration object used to communicate with loop body function.
| 972 | |
| 973 | |
| 974 | class PForConfig(object): |
| 975 | """A configuration object used to communicate with loop body function.""" |
| 976 | |
| 977 | def __init__(self): |
| 978 | # This may be set to the number of iterations. |
| 979 | self._maybe_iters = None |
| 980 | # Map from output placeholder to the unvectorized tensor. |
| 981 | self._reduce_concat_map = object_identity.ObjectIdentityDictionary() |
| 982 | # Reverse map of `self._reduce_concat_map`. |
| 983 | self._reverse_reduce_concat_map = object_identity.ObjectIdentityDictionary() |
| 984 | |
| 985 | def _has_reductions(self): |
| 986 | """True if some reductions where performed by loop body.""" |
| 987 | return len(self._reduce_concat_map) |
| 988 | |
| 989 | def _set_iters(self, iters): |
| 990 | """Set number of pfor iterations.""" |
| 991 | self._maybe_iters = iters |
| 992 | |
| 993 | # TODO(agarwal): handle reductions inside control flow constructs. |
| 994 | def reduce_concat(self, x): |
| 995 | """Performs a concat reduction on `x` across pfor iterations. |
| 996 | |
| 997 | Note that this currently may not work inside a control flow construct. |
| 998 | Args: |
| 999 | x: an unvectorized Tensor. |
| 1000 | |
| 1001 | Returns: |
| 1002 | A Tensor that has rank one higher than `x`. The value is the vectorized |
| 1003 | version of `x`, i.e. stacking the value of `x` across different pfor |
| 1004 | iterations. |
| 1005 | """ |
| 1006 | assert not context.executing_eagerly() |
| 1007 | assert isinstance(x, ops.Tensor) |
| 1008 | if x not in self._reduce_concat_map: |
| 1009 | out_shape = tensor_shape.TensorShape([self._maybe_iters]).concatenate( |
| 1010 | x.shape) |
| 1011 | with ops.control_dependencies([x]): |
| 1012 | # Control dependency to make sure out is converted after x. |
| 1013 | out = array_ops.placeholder(x.dtype, out_shape) |
| 1014 | self._reduce_concat_map[out] = x |
| 1015 | self._reverse_reduce_concat_map[x] = out |
| 1016 | return out |
| 1017 | else: |
| 1018 | return self._reverse_reduce_concat_map[x] |
| 1019 | |
| 1020 | def reduce_mean(self, x): |
| 1021 | """Performs a mean reduction on `x` across pfor iterations. |
| 1022 | |
| 1023 | Note that this currently may not work inside a control flow construct. |
| 1024 | Args: |
| 1025 | x: an unvectorized Tensor. |
| 1026 | |
| 1027 | Returns: |
| 1028 | A Tensor that has same rank as `x`. The value is the mean of the values |
| 1029 | of `x` across the pfor iterations. |
| 1030 | """ |
| 1031 | y = self.reduce_concat(x) |