Helper method to perform transpose and reshape for contraction op. This method is helpful in reducing `math_ops.tensordot` to `math_ops.matmul` using `array_ops.transpose` and `array_ops.reshape`. The method takes a tensor and performs the correct transpose and reshape operation for a g
(a, axes, flipped=False)
| 3962 | """ |
| 3963 | |
| 3964 | def _tensordot_reshape(a, axes, flipped=False): |
| 3965 | """Helper method to perform transpose and reshape for contraction op. |
| 3966 | |
| 3967 | This method is helpful in reducing `math_ops.tensordot` to `math_ops.matmul` |
| 3968 | using `array_ops.transpose` and `array_ops.reshape`. The method takes a |
| 3969 | tensor and performs the correct transpose and reshape operation for a given |
| 3970 | set of indices. It returns the reshaped tensor as well as a list of indices |
| 3971 | necessary to reshape the tensor again after matrix multiplication. |
| 3972 | |
| 3973 | Args: |
| 3974 | a: `Tensor`. |
| 3975 | axes: List or `int32` `Tensor` of unique indices specifying valid axes of |
| 3976 | `a`. |
| 3977 | flipped: An optional `bool`. Defaults to `False`. If `True`, the method |
| 3978 | assumes that `a` is the second argument in the contraction operation. |
| 3979 | |
| 3980 | Returns: |
| 3981 | A tuple `(reshaped_a, free_dims, free_dims_static)` where `reshaped_a` is |
| 3982 | the tensor `a` reshaped to allow contraction via `matmul`, `free_dims` is |
| 3983 | either a list of integers or an `int32` `Tensor`, depending on whether |
| 3984 | the shape of a is fully specified, and free_dims_static is either a list |
| 3985 | of integers and None values, or None, representing the inferred |
| 3986 | static shape of the free dimensions |
| 3987 | """ |
| 3988 | if a.get_shape().is_fully_defined() and isinstance(axes, (list, tuple)): |
| 3989 | shape_a = a.get_shape().as_list() |
| 3990 | axes = [i if i >= 0 else i + len(shape_a) for i in axes] |
| 3991 | free = [i for i in xrange(len(shape_a)) if i not in axes] |
| 3992 | free_dims = [shape_a[i] for i in free] |
| 3993 | prod_free = int(np.prod([shape_a[i] for i in free])) |
| 3994 | prod_axes = int(np.prod([shape_a[i] for i in axes])) |
| 3995 | perm = list(axes) + free if flipped else free + list(axes) |
| 3996 | new_shape = [prod_axes, prod_free] if flipped else [prod_free, prod_axes] |
| 3997 | reshaped_a = array_ops.reshape(array_ops.transpose(a, perm), new_shape) |
| 3998 | return reshaped_a, free_dims, free_dims |
| 3999 | else: |
| 4000 | if a.get_shape().ndims is not None and isinstance(axes, (list, tuple)): |
| 4001 | shape_a = a.get_shape().as_list() |
| 4002 | axes = [i if i >= 0 else i + len(shape_a) for i in axes] |
| 4003 | free = [i for i in xrange(len(shape_a)) if i not in axes] |
| 4004 | axes_dims = [shape_a[i] for i in axes] |
| 4005 | free_dims = [shape_a[i] for i in free] |
| 4006 | free_dims_static = free_dims |
| 4007 | axes = ops.convert_to_tensor(axes, dtype=dtypes.int32, name="axes") |
| 4008 | free = ops.convert_to_tensor(free, dtype=dtypes.int32, name="free") |
| 4009 | shape_a = array_ops.shape(a) |
| 4010 | else: |
| 4011 | free_dims_static = None |
| 4012 | shape_a = array_ops.shape(a) |
| 4013 | rank_a = array_ops.rank(a) |
| 4014 | axes = ops.convert_to_tensor(axes, dtype=dtypes.int32, name="axes") |
| 4015 | axes = array_ops.where(axes >= 0, axes, axes + rank_a) |
| 4016 | free, _ = array_ops.setdiff1d(range(rank_a), axes) |
| 4017 | free_dims = array_ops.gather(shape_a, free) |
| 4018 | axes_dims = array_ops.gather(shape_a, axes) |
| 4019 | prod_free_dims = reduce_prod(free_dims) |
| 4020 | prod_axes_dims = reduce_prod(axes_dims) |
| 4021 | if flipped: |
no test coverage detected