Ragged compatible squeeze. If `input` is a `tf.Tensor`, then this calls `tf.squeeze`. If `input` is a `tf.RaggedTensor`, then this operation takes `O(N)` time, where `N` is the number of elements in the squeezed dimensions. Args: input: A potentially ragged tensor. The input to squeez
(input, axis=None, name=None)
| 30 | |
| 31 | |
| 32 | def squeeze(input, axis=None, name=None): # pylint: disable=redefined-builtin |
| 33 | """Ragged compatible squeeze. |
| 34 | |
| 35 | If `input` is a `tf.Tensor`, then this calls `tf.squeeze`. |
| 36 | |
| 37 | If `input` is a `tf.RaggedTensor`, then this operation takes `O(N)` time, |
| 38 | where `N` is the number of elements in the squeezed dimensions. |
| 39 | |
| 40 | Args: |
| 41 | input: A potentially ragged tensor. The input to squeeze. |
| 42 | axis: An optional list of ints. Defaults to `None`. If the `input` is |
| 43 | ragged, it only squeezes the dimensions listed. It fails if `input` is |
| 44 | ragged and axis is []. If `input` is not ragged it calls tf.squeeze. Note |
| 45 | that it is an error to squeeze a dimension that is not 1. It must be in |
| 46 | the range of [-rank(input), rank(input)). |
| 47 | name: A name for the operation (optional). |
| 48 | |
| 49 | Returns: |
| 50 | A potentially ragged tensor. Contains the same data as input, |
| 51 | but has one or more dimensions of size 1 removed. |
| 52 | """ |
| 53 | with ops.name_scope(name, 'RaggedSqueeze', [input]): |
| 54 | input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input) |
| 55 | if isinstance(input, ops.Tensor): |
| 56 | return array_ops.squeeze(input, axis, name) |
| 57 | |
| 58 | if axis is None: |
| 59 | raise ValueError('Ragged.squeeze must have an axis argument.') |
| 60 | if isinstance(axis, int): |
| 61 | axis = [axis] |
| 62 | elif ((not isinstance(axis, (list, tuple))) or |
| 63 | (not all(isinstance(d, int) for d in axis))): |
| 64 | raise TypeError('Axis must be a list or tuple of integers.') |
| 65 | |
| 66 | dense_dims = [] |
| 67 | ragged_dims = [] |
| 68 | # Normalize all the dims in axis to be positive |
| 69 | axis = [ragged_util.get_positive_axis(d, input.shape.ndims) for d in axis] |
| 70 | for dim in axis: |
| 71 | if dim > input.ragged_rank: |
| 72 | dense_dims.append(dim - input.ragged_rank) |
| 73 | else: |
| 74 | ragged_dims.append(dim) |
| 75 | |
| 76 | # Make sure the specified ragged dimensions are squeezable. |
| 77 | assertion_list = [] |
| 78 | scalar_tensor_one = constant_op.constant(1, dtype=input.row_splits.dtype) |
| 79 | for i, r in enumerate(input.nested_row_lengths()): |
| 80 | if i + 1 in ragged_dims: |
| 81 | assertion_list.append( |
| 82 | control_flow_ops.Assert( |
| 83 | math_ops.reduce_all(math_ops.equal(r, scalar_tensor_one)), |
| 84 | ['the given axis (axis = %d) is not squeezable!' % (i + 1)])) |
| 85 | if 0 in ragged_dims: |
| 86 | scalar_tensor_two = constant_op.constant(2, dtype=dtypes.int32) |
| 87 | assertion_list.append( |
| 88 | control_flow_ops.Assert( |
| 89 | math_ops.equal( |
nothing calls this directly
no test coverage detected