Helper class to make it easier to manage caching in `Bijector`.
| 43 | |
| 44 | |
| 45 | class _Mapping(collections.namedtuple( |
| 46 | "_Mapping", ["x", "y", "ildj_map", "kwargs"])): |
| 47 | """Helper class to make it easier to manage caching in `Bijector`.""" |
| 48 | |
| 49 | def __new__(cls, x=None, y=None, ildj_map=None, kwargs=None): |
| 50 | """Custom __new__ so namedtuple items have defaults. |
| 51 | |
| 52 | Args: |
| 53 | x: `Tensor`. Forward. |
| 54 | y: `Tensor`. Inverse. |
| 55 | ildj_map: `Dictionary`. This is a mapping from event_ndims to a `Tensor` |
| 56 | representing the inverse log det jacobian. |
| 57 | kwargs: Python dictionary. Extra args supplied to |
| 58 | forward/inverse/etc functions. |
| 59 | |
| 60 | Returns: |
| 61 | mapping: New instance of _Mapping. |
| 62 | """ |
| 63 | return super(_Mapping, cls).__new__(cls, x, y, ildj_map, kwargs) |
| 64 | |
| 65 | @property |
| 66 | def x_key(self): |
| 67 | """Returns key used for caching Y=g(X).""" |
| 68 | return ((object_identity.Reference(self.x),) + |
| 69 | self._deep_tuple(tuple(sorted(self.kwargs.items())))) |
| 70 | |
| 71 | @property |
| 72 | def y_key(self): |
| 73 | """Returns key used for caching X=g^{-1}(Y).""" |
| 74 | return ((object_identity.Reference(self.y),) + |
| 75 | self._deep_tuple(tuple(sorted(self.kwargs.items())))) |
| 76 | |
| 77 | def merge(self, x=None, y=None, ildj_map=None, kwargs=None, mapping=None): |
| 78 | """Returns new _Mapping with args merged with self. |
| 79 | |
| 80 | Args: |
| 81 | x: `Tensor`. Forward. |
| 82 | y: `Tensor`. Inverse. |
| 83 | ildj_map: `Dictionary`. This is a mapping from event_ndims to a `Tensor` |
| 84 | representing the inverse log det jacobian. |
| 85 | kwargs: Python dictionary. Extra args supplied to |
| 86 | forward/inverse/etc functions. |
| 87 | mapping: Instance of _Mapping to merge. Can only be specified if no other |
| 88 | arg is specified. |
| 89 | |
| 90 | Returns: |
| 91 | mapping: New instance of `_Mapping` which has inputs merged with self. |
| 92 | |
| 93 | Raises: |
| 94 | ValueError: if mapping and any other arg is not `None`. |
| 95 | """ |
| 96 | if mapping is None: |
| 97 | mapping = _Mapping(x=x, y=y, ildj_map=ildj_map, kwargs=kwargs) |
| 98 | elif any(arg is not None for arg in [x, y, ildj_map, kwargs]): |
| 99 | raise ValueError("Cannot simultaneously specify mapping and individual " |
| 100 | "arguments.") |
| 101 | |
| 102 | return _Mapping( |