def _get_dataset(): inputs = array_ops.expand_dims_v2(constant_op.constant(range(10)), axis=1) targets = array_ops.expand_dims_v2( constant_op.constant(range(10)), axis=1) # Make global batch size 12 for 2 replicas and a non-repeated dataset with # 10 elements so that we have partial batch dataset = dataset_ops.Dataset.from_tensor_slices( (inputs, targets)).batch(12, drop_remainder=False) return dataset
def match_dtype_and_rank(y_t, y_p, sw): """Match dtype and rank of predictions.""" if y_t.shape.rank == 1 and y_p.shape.rank == 2: y_t = array_ops.expand_dims_v2(y_t, axis=-1) if sw is not None: if sw.shape.rank == 1 and y_p.shape.rank == 2: sw = array_ops.expand_dims_v2(sw, axis=-1) # Dtype. y_t = math_ops.cast(y_t, y_p.dtype) if sw is not None: sw = math_ops.cast(sw, y_p.dtype) return y_t, y_p, sw
def _gather(strategy, value): """Gathers a single value.""" # pylint: disable=protected-access if not isinstance(value, values.DistributedValues): assert isinstance(value, core.Tensor) value = values.PerReplica([value]) if not isinstance( strategy.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended): return array_ops.stack(value._values) assert len(strategy.extended.worker_devices) == len(value._values) inputs = [array_ops.expand_dims_v2(v, axis=0) for v in value._values] collective_keys = strategy.extended._collective_keys devices = strategy.extended.worker_devices group_size = strategy.num_replicas_in_sync @def_function.function def gather_fn(): gathered = cross_device_utils.build_collective_gather( inputs, devices, group_size, collective_keys) return distribute_utils.update_regroup(strategy.extended, gathered, group=True) return gather_fn()
def call(self, inputs): if self._channels_first: rank = inputs.shape.rank if rank and rank > 1: # Switch to channels-last format. permutation = [0] permutation.extend(range(2, rank)) permutation.append(1) inputs = array_ops.transpose(inputs, perm=permutation) if context.executing_eagerly(): # Full static shape is guaranteed to be available. # Performance: Using `constant_op` is much faster than passing a list. flattened_shape = constant_op.constant([inputs.shape[0], -1]) return gen_array_ops.reshape(inputs, flattened_shape) else: input_shape = inputs.shape rank = input_shape.rank if rank == 1: return array_ops.expand_dims_v2(inputs, axis=1) else: batch_dim = tensor_shape.dimension_value(input_shape[0]) non_batch_dims = input_shape[1:] # Reshape in a way that preserves as much shape info as possible. if non_batch_dims.is_fully_defined(): last_dim = int( functools.reduce(operator.mul, non_batch_dims)) flattened_shape = constant_op.constant([-1, last_dim]) elif batch_dim is not None: flattened_shape = constant_op.constant( [int(batch_dim), -1]) else: flattened_shape = [array_ops.shape_v2(inputs)[0], -1] return array_ops.reshape(inputs, flattened_shape)
def _partial_shaped_bools(): rand_vect = math_ops.range( random_ops.random_uniform(shape=(), minval=2, maxval=3, dtype=dtypes.int32)) return array_ops.expand_dims_v2(rand_vect, 0) < 0
def match_dtype_and_rank(y_t, y_p, sw): """Match dtype and rank of predictions.""" if y_t.shape.rank == 1 and y_p.shape.rank == 2: y_t = array_ops.expand_dims_v2(y_t, axis=-1) if sw is not None: if sw.shape.rank == 1 and y_p.shape.rank == 2: sw = array_ops.expand_dims_v2(sw, axis=-1) # Dtype. # This is required mainly for custom loss functions which do not take care # casting dtypes. if ((y_t.dtype.is_floating and y_p.dtype.is_floating) or (y_t.dtype.is_integer and y_p.dtype.is_integer)): y_t = math_ops.cast(y_t, y_p.dtype) if sw is not None: sw = math_ops.cast(sw, y_p.dtype) return y_t, y_p, sw
def _get_total_loss_tensor(activations): losses = [] for activation in activations: losses.append( math_ops.reduce_mean( math_ops.reduce_sum( gen_math_ops.squared_difference(activation, 0), 1))) total_loss = array_ops.expand_dims_v2(sum(losses), 0) return total_loss
def _gather(strategy, value): """Gathers a single value.""" # pylint: disable=protected-access if not isinstance(value, values.DistributedValues): value = values.PerReplica([ops.convert_to_tensor(value)]) if not isinstance(strategy.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended): return array_ops.stack(value._values) assert len(strategy.extended.worker_devices) == len(value._values) inputs = [array_ops.expand_dims_v2(v, axis=0) for v in value._values] return strategy.gather(values.PerReplica(inputs), axis=0)
def all_gather( self, input_tensor: core.TensorLike, axis: core.TensorLike, options: Optional[collective_util.Options] = None) -> core.Tensor: """All-gather a dense tensor. This method must be called inside a tf.function. Args: input_tensor: a dense tensor. It must have the same rank on all replicas, and dimensions other than `axis` need to be the same as well. axis: 0-D int32 Tensor. Dimension along which to gather. Must be in the range [0, rank(value)). options: an optional tf.distribute.experimental.CommunicationOptions. If provided, it overrides the default options. Returns: The gathered Tensor. Raises: RuntimeError: if called in eager mode. """ if context.executing_eagerly(): raise RuntimeError('all_gather is not supported in eager mode.') with ops.device(self._device), \ ops.control_dependencies([array_ops.identity(input_tensor)]): # 1. Transpose # E.g. Given an input_tensor with shape [2,2,5,1] and axis to gather is 3, # we use perm_pre=[3 0 1 2] to reshape it to [1,2,2,5], which # brings the 3rd dim first; afterwards we use perm_after=[1,2,3,0] to # place it back. perm_pre = array_ops.concat( ([axis], math_ops.range(axis), math_ops.range(axis + 1, array_ops.rank(input_tensor))), axis=0) input_tensor_t = array_ops.transpose(input_tensor, perm=perm_pre) # 2. Pad gathered_shape = self._all_gather( array_ops.expand_dims_v2(array_ops.shape_v2(input_tensor_t), axis=0), options) first_dims = gathered_shape[:, 0] full_axis_dim = math_ops.reduce_max(first_dims) padded_input_tensor = _pad_util(input_tensor_t, full_axis_dim) # 3. Gather gather_padded_out_tensor = self._all_gather( padded_input_tensor, options) # 4. Unpad split_tensors = [] for i in range(self._group_size): start_pos = i * full_axis_dim split_tensors.append( gather_padded_out_tensor[start_pos:start_pos + first_dims[i]]) out_tensor_t = array_ops.concat(split_tensors, 0) # 5. Transpose back perm_after = array_ops.concat( (math_ops.range(1, axis + 1), [0], math_ops.range(axis + 1, array_ops.rank(input_tensor_t))), axis=0) return array_ops.transpose(out_tensor_t, perm=perm_after)
def pinv(a, rcond=None, validate_args=False, name=None): """Compute the Moore-Penrose pseudo-inverse of one or more matrices. Calculate the [generalized inverse of a matrix]( https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse) using its singular-value decomposition (SVD) and including all large singular values. The pseudo-inverse of a matrix `A`, is defined as: 'the matrix that 'solves' [the least-squares problem] `A @ x = b`,' i.e., if `x_hat` is a solution, then `A_pinv` is the matrix such that `x_hat = A_pinv @ b`. It can be shown that if `U @ Sigma @ V.T = A` is the singular value decomposition of `A`, then `A_pinv = V @ inv(Sigma) U^T`. [(Strang, 1980)][1] This function is analogous to [`numpy.linalg.pinv`]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.pinv.html). It differs only in default value of `rcond`. In `numpy.linalg.pinv`, the default `rcond` is `1e-15`. Here the default is `10. * max(num_rows, num_cols) * np.finfo(dtype).eps`. Args: a: (Batch of) `float`-like matrix-shaped `Tensor`(s) which are to be pseudo-inverted. rcond: `Tensor` of small singular value cutoffs. Singular values smaller (in modulus) than `rcond` * largest_singular_value (again, in modulus) are set to zero. Must broadcast against `tf.shape(a)[:-2]`. Default value: `10. * max(num_rows, num_cols) * np.finfo(a.dtype).eps`. validate_args: When `True`, additional assertions might be embedded in the graph. Default value: `False` (i.e., no graph assertions are added). name: Python `str` prefixed to ops created by this function. Default value: 'pinv'. Returns: a_pinv: (Batch of) pseudo-inverse of input `a`. Has same shape as `a` except rightmost two dimensions are transposed. Raises: TypeError: if input `a` does not have `float`-like `dtype`. ValueError: if input `a` has fewer than 2 dimensions. #### Examples ```python import tensorflow as tf import tensorflow_probability as tfp a = tf.constant([[1., 0.4, 0.5], [0.4, 0.2, 0.25], [0.5, 0.25, 0.35]]) tf.matmul(tf.linalg..pinv(a), a) # ==> array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]], dtype=float32) a = tf.constant([[1., 0.4, 0.5, 1.], [0.4, 0.2, 0.25, 2.], [0.5, 0.25, 0.35, 3.]]) tf.matmul(tf.linalg..pinv(a), a) # ==> array([[ 0.76, 0.37, 0.21, -0.02], [ 0.37, 0.43, -0.33, 0.02], [ 0.21, -0.33, 0.81, 0.01], [-0.02, 0.02, 0.01, 1. ]], dtype=float32) ``` #### References [1]: G. Strang. 'Linear Algebra and Its Applications, 2nd Ed.' Academic Press, Inc., 1980, pp. 139-142. """ with ops.name_scope(name or 'pinv'): a = ops.convert_to_tensor(a, name='a') assertions = _maybe_validate_matrix(a, validate_args) if assertions: with ops.control_dependencies(assertions): a = array_ops.identity(a) dtype = a.dtype.as_numpy_dtype if rcond is None: def get_dim_size(dim): dim_val = tensor_shape.dimension_value(a.shape[dim]) if dim_val is not None: return dim_val return array_ops.shape(a)[dim] num_rows = get_dim_size(-2) num_cols = get_dim_size(-1) if isinstance(num_rows, int) and isinstance(num_cols, int): max_rows_cols = float(max(num_rows, num_cols)) else: max_rows_cols = math_ops.cast( math_ops.maximum(num_rows, num_cols), dtype) rcond = 10. * max_rows_cols * np.finfo(dtype).eps rcond = ops.convert_to_tensor(rcond, dtype=dtype, name='rcond') # Calculate pseudo inverse via SVD. # Note: if a is Hermitian then u == v. (We might observe additional # performance by explicitly setting `v = u` in such cases.) [ singular_values, # Sigma left_singular_vectors, # U right_singular_vectors, # V ] = svd(a, full_matrices=False, compute_uv=True) # Saturate small singular values to inf. This has the effect of make # `1. / s = 0.` while not resulting in `NaN` gradients. cutoff = rcond * math_ops.reduce_max(singular_values, axis=-1) singular_values = array_ops.where_v2( singular_values > array_ops.expand_dims_v2(cutoff, -1), singular_values, np.array(np.inf, dtype)) # By the definition of the SVD, `a == u @ s @ v^H`, and the pseudo-inverse # is defined as `pinv(a) == v @ inv(s) @ u^H`. a_pinv = math_ops.matmul(right_singular_vectors / array_ops.expand_dims_v2(singular_values, -2), left_singular_vectors, adjoint_b=True) if a.shape is not None and a.shape.rank is not None: a_pinv.set_shape(a.shape[:-2].concatenate( [a.shape[-1], a.shape[-2]])) return a_pinv
def _expand_single_1d_tensor(t): # Leaves `CompositeTensor`s as-is. if (isinstance(t, ops.Tensor) and isinstance(t.shape, tensor_shape.TensorShape) and t.shape.rank == 1): return array_ops.expand_dims_v2(t, axis=-1) return t
def audio(name, tensor, sample_rate, max_outputs=3, collections=None, family=None): # pylint: disable=line-too-long """Outputs a `Summary` protocol buffer with audio. The summary has up to `max_outputs` summary values containing audio. The audio is built from `tensor` which must be 3-D with shape `[batch_size, frames, channels]` or 2-D with shape `[batch_size, frames]`. The values are assumed to be in the range of `[-1.0, 1.0]` with a sample rate of `sample_rate`. The `tag` in the outputted Summary.Value protobufs is generated based on the name, with a suffix depending on the max_outputs setting: * If `max_outputs` is 1, the summary value tag is '*name*/audio'. * If `max_outputs` is greater than 1, the summary value tags are generated sequentially as '*name*/audio/0', '*name*/audio/1', etc Args: name: A name for the generated node. Will also serve as a series name in TensorBoard. tensor: A 3-D `float32` `Tensor` of shape `[batch_size, frames, channels]` or a 2-D `float32` `Tensor` of shape `[batch_size, frames]`. sample_rate: A Scalar `float32` `Tensor` indicating the sample rate of the signal in hertz. max_outputs: Max number of batch elements to generate audio for. collections: Optional list of ops.GraphKeys. The collections to add the summary to. Defaults to [_ops.GraphKeys.SUMMARIES] family: Optional; if provided, used as the prefix of the summary tag name, which controls the tab name used for display on Tensorboard. Returns: A scalar `Tensor` of type `string`. The serialized `Summary` protocol buffer. @compatibility(TF2) For compatibility purposes, when invoked in TF2 where the outermost context is eager mode, this API will check if there is a suitable TF2 summary writer context available, and if so will forward this call to that writer instead. A "suitable" writer context means that the writer is set as the default writer, and there is an associated non-empty value for `step` (see `tf.summary.SummaryWriter.as_default`, or alternatively `tf.summary.experimental.set_step`). For the forwarded call, the arguments here will be passed to the TF2 implementation of `tf.summary.audio`, and the return value will be an empty bytestring tensor, to avoid duplicate summary writing. This forwarding is best-effort and not all arguments will be preserved. Additionally: * The TF2 op just outputs the data under a single tag that contains multiple samples, rather than multiple tags (i.e. no "/0" or "/1" suffixes). To migrate to TF2, please use `tf.summary.audio` instead. Please check [Migrating tf.summary usage to TF 2.0](https://www.tensorflow.org/tensorboard/migrate#in_tf_1x) for concrete steps for migration. #### How to Map Arguments | TF1 Arg Name | TF2 Arg Name | Note | | :------------ | :-------------- | :------------------------------------- | | `name` | `name` | - | | `tensor` | `data` | Input for this argument now must be | : : : three-dimensional `[k, t, c]`, where : : : : `k` is the number of audio clips, `t` : : : : is the number of frames, and `c` is : : : : the number of channels. Two-dimensional: : : : input is no longer supported. : | `sample_rate` | `sample_rate` | - | | - | `step` | Explicit int64-castable monotonic step | : : : value. If omitted, this defaults to : : : : `tf.summary.experimental.get_step()`. : | `max_outputs` | `max_outputs` | - | | `collections` | Not Supported | - | | `family` | Removed | Please use `tf.name_scope` instead to | : : : manage summary name prefix. : | - | `encoding` | Optional constant str for the desired | : : : encoding. Check the docs for : : : : `tf.summary.audio` for latest supported: : : : audio formats. : | - | `description` | Optional long-form `str` description | : : : for the summary. Markdown is supported.: : : : Defaults to empty. : @end_compatibility """ # Special case: invoke v2 op for TF2 users who have a v2 writer. if _should_invoke_v2_op(): # Defer the import to happen inside the symbol to prevent breakage due to # missing dependency. from tensorboard.summary.v2 import audio as audio_v2 # pylint: disable=g-import-not-at-top if tensor.shape.rank == 2: # TF2 op requires 3-D tensor, add the `channels` dimension. tensor = _array_ops.expand_dims_v2(tensor, axis=2) with _compat_summary_scope(name, family) as tag: audio_v2( name=tag, data=tensor, sample_rate=sample_rate, step=_get_step_for_v2(), max_outputs=max_outputs, ) # Return an empty Tensor, which will be acceptable as an input to the # `tf.compat.v1.summary.merge()` API. return _constant_op.constant(b'') # Fall back to legacy v1 audio implementation. if _distribute_summary_op_util.skip_summary(): return _constant_op.constant('') with _summary_op_util.summary_scope( name, family=family, values=[tensor]) as (tag, scope): sample_rate = _ops.convert_to_tensor( sample_rate, dtype=_dtypes.float32, name='sample_rate') val = _gen_logging_ops.audio_summary_v2( tag=tag, tensor=tensor, max_outputs=max_outputs, sample_rate=sample_rate, name=scope) _summary_op_util.collect(val, collections, [_ops.GraphKeys.SUMMARIES]) return val
def build_collective_gather(input_tensors, devices, group_size, collective_keys, axis, communication_hint='AUTO', control_inputs=None, timeout=None): """Build a subgraph that does one full all-gather, using the collective Op. This method must be called in graph mode or inside a tf.function. Args: input_tensors: tensors within a single worker graph that are to be gathered together; must be one per device. Input tensors cannot have rank 0. devices: a list of device strings to run the collective on. group_size: total number of devices globally that will be doing this same gathering. The gathering will actually include the corresponding tensors at all these workers. collective_keys: a CollectiveKeys object. axis: 0-D int32 Tensor. Dimension along which to gather. Must be in the range [0, rank(value)). communication_hint: string providing hint to runtime for choosing collective implementation. Available options are `AUTO`, `NCCL`, and `RING`. control_inputs: if not None, add control edges between control_inputs and (index-wise) corresponding collective_gather tensors timeout: a float or None. The timeout in seconds. Returns: An array of final tensors, one per device, computed by the full gather. """ if len(input_tensors) != len(devices): raise ValueError( 'collective requires one input tensor for each device, %d != %d' % (len(input_tensors), len(devices))) if group_size < 2: return input_tensors group_key = collective_keys.get_group_key(devices) instance_key_tensor = collective_keys.get_op_instance_key() instance_key_shape = collective_keys.get_op_instance_key() out_tensors = [] for idx, input_tensor in enumerate(input_tensors): with ops.device(devices[idx]), ops.control_dependencies( _control_input(devices, control_inputs, idx)): # 1. Transpose # E.g. Given an input_tensor with shape [2,2,5,1] and axis to gather is 3, # we use perm_pre=[3 0 1 2] to reshape it to [1,2,2,5], which # brings the 3rd dim first; afterwards we use perm_after=[1,2,3,0] to # place it back. perm_pre = array_ops.concat( ([axis], math_ops.range(axis), math_ops.range(axis + 1, array_ops.rank(input_tensor))), axis=0) input_tensor_t = array_ops.transpose(input_tensor, perm=perm_pre) # 2. Pad gathered_shape = collective_ops.all_gather( array_ops.expand_dims_v2(array_ops.shape_v2(input_tensor_t), axis=0), group_size, group_key, instance_key_shape, communication_hint, timeout=timeout) first_dims = gathered_shape[:, 0] full_axis_dim = math_ops.reduce_max(first_dims) padded_input_tensor = _pad_util(input_tensor_t, full_axis_dim) # 3. Gather gather_padded_out_tensor = collective_ops.all_gather( padded_input_tensor, group_size, group_key, instance_key_tensor, communication_hint, timeout=timeout) # 4. Unpad split_tensors = [] for i in range(first_dims.shape[0]): start_pos = i * full_axis_dim split_tensors.append(gather_padded_out_tensor[start_pos:start_pos + first_dims[i]]) out_tensor_t = array_ops.concat(split_tensors, 0) # 5. Transpose back perm_after = array_ops.concat( (math_ops.range(1, axis + 1), [0], math_ops.range(axis + 1, array_ops.rank(input_tensor_t))), axis=0) out_tensor = array_ops.transpose(out_tensor_t, perm=perm_after) out_tensors.append(out_tensor) return out_tensors
def all_gather(self, input_tensor, axis, communication_hint='AUTO', timeout=0): """All-gather a dense tensor. This method must be called inside a tf.function. Args: input_tensor: a dense tensor. It must have the same rank on all replicas, and dimensions other than `axis` need to be the same as well. axis: 0-D int32 Tensor. Dimension along which to gather. Must be in the range [0, rank(value)). communication_hint: string providing hint to runtime for choosing collective implementation. Available options are `AUTO`, `NCCL`, and `RING`. timeout: a float. The timeout in seconds. Returns: The gathered Tensor. Raises: RuntimeError: if called in eager mode. """ if context.executing_eagerly(): raise RuntimeError('all_gather in eager mode is not supported') with ops.device(self._device), \ ops.control_dependencies([array_ops.identity(input_tensor)]): # 1. Transpose # E.g. Given an input_tensor with shape [2,2,5,1] and axis to gather is 3, # we use perm_pre=[3 0 1 2] to reshape it to [1,2,2,5], which # brings the 3rd dim first; afterwards we use perm_after=[1,2,3,0] to # place it back. perm_pre = array_ops.concat( ([axis], math_ops.range(axis), math_ops.range(axis + 1, array_ops.rank(input_tensor))), axis=0) input_tensor_t = array_ops.transpose(input_tensor, perm=perm_pre) # 2. Pad gathered_shape = self._all_gather(array_ops.expand_dims_v2( array_ops.shape_v2(input_tensor_t), axis=0), communication_hint, timeout=timeout) first_dims = gathered_shape[:, 0] full_axis_dim = math_ops.reduce_max(first_dims) padded_input_tensor = _pad_util(input_tensor_t, full_axis_dim) # 3. Gather gather_padded_out_tensor = self._all_gather(padded_input_tensor, communication_hint, timeout=timeout) # 4. Unpad split_tensors = [] for i in range(self._group_size): start_pos = i * full_axis_dim split_tensors.append( gather_padded_out_tensor[start_pos:start_pos + first_dims[i]]) out_tensor_t = array_ops.concat(split_tensors, 0) # 5. Transpose back perm_after = array_ops.concat( (math_ops.range(1, axis + 1), [0], math_ops.range(axis + 1, array_ops.rank(input_tensor_t))), axis=0) return array_ops.transpose(out_tensor_t, perm=perm_after)
def pinv(a, rcond=None, validate_args=False, name=None): """Compute the Moore-Penrose pseudo-inverse of one or more matrices.""" with ops.name_scope(name or 'pinv'): a = ops.convert_to_tensor(a, name='a') assertions = _maybe_validate_matrix(a, validate_args) if assertions: with ops.control_dependencies(assertions): a = array_ops.identity(a) dtype = a.dtype.as_numpy_dtype if rcond is None: def get_dim_size(dim): dim_val = tensor_shape.dimension_value(a.shape[dim]) if dim_val is not None: return dim_val return array_ops.shape(a)[dim] num_rows = get_dim_size(-2) num_cols = get_dim_size(-1) if isinstance(num_rows, int) and isinstance(num_cols, int): max_rows_cols = float(max(num_rows, num_cols)) else: max_rows_cols = math_ops.cast( math_ops.maximum(num_rows, num_cols), dtype) rcond = 10. * max_rows_cols * np.finfo(dtype).eps rcond = ops.convert_to_tensor(rcond, dtype=dtype, name='rcond') # Calculate pseudo inverse via SVD. # Note: if a is Hermitian then u == v. (We might observe additional # performance by explicitly setting `v = u` in such cases.) [ singular_values, # Sigma left_singular_vectors, # U right_singular_vectors, # V ] = tf.svd(a, full_matrices=False, compute_uv=True) # Saturate small singular values to inf. This has the effect of make # `1. / s = 0.` while not resulting in `NaN` gradients. cutoff = rcond * math_ops.reduce_max(singular_values, axis=-1) singular_values = array_ops.where( singular_values > array_ops.expand_dims_v2(cutoff, -1), singular_values, np.array(np.inf, dtype).repeat(singular_values.get_shape().as_list()[0])) # By the definition of the SVD, `a == u @ s @ v^H`, and the pseudo-inverse # is defined as `pinv(a) == v @ inv(s) @ u^H`. a_pinv = math_ops.matmul(right_singular_vectors / array_ops.expand_dims_v2(singular_values, -2), left_singular_vectors, adjoint_b=True) if a.shape is not None and a.shape.rank is not None: a_pinv.set_shape(a.shape[:-2].concatenate( [a.shape[-1], a.shape[-2]])) return a_pinv
def test_slicing(self): v = [ variables_lib.Variable([[1, 2], [3, 4], [5, 6]]), variables_lib.Variable([[7, 8], [9, 10], [11, 12]]), variables_lib.Variable([[13, 14], [15, 16]]) ] sv = sharded_variable.ShardedVariable(v) empty = v[0][0:0] # Test cases: positive step self.assertAllEqual(sv[:], array_ops.concat(v, axis=0)) self.assertAllEqual(sv[:2], [[1, 2], [3, 4]]) self.assertAllEqual(sv[-8:2], [[1, 2], [3, 4]]) self.assertAllEqual(sv[-10:2], [[1, 2], [3, 4]]) self.assertAllEqual(sv[5:], [[11, 12], [13, 14], [15, 16]]) self.assertAllEqual(sv[5:-1], [[11, 12], [13, 14]]) self.assertAllEqual(sv[::3], [[1, 2], [7, 8], [13, 14]]) self.assertAllEqual(sv[::5], [[1, 2], [11, 12]]) self.assertAllEqual(sv[1::6], [[3, 4], [15, 16]]) self.assertAllEqual(sv[1:5:6], [[3, 4]]) self.assertAllEqual(sv[1::7], [[3, 4]]) self.assertAllEqual(sv[2:7], [[5, 6], [7, 8], [9, 10], [11, 12], [13, 14]]) self.assertAllEqual(sv[2:7:2], [[5, 6], [9, 10], [13, 14]]) self.assertAllEqual(sv[2:7:3], [[5, 6], [11, 12]]) # Test cases: negative step self.assertAllEqual( sv[::-1], array_ops.reverse(array_ops.concat(v, axis=0), axis=[0])) self.assertAllEqual(sv[2::-1], [[5, 6], [3, 4], [1, 2]]) self.assertAllEqual(sv[2:-8:-1], [[5, 6], [3, 4]]) self.assertAllEqual(sv[2:-10:-1], [[5, 6], [3, 4], [1, 2]]) self.assertAllEqual(sv[4::-1], [[9, 10], [7, 8], [5, 6], [3, 4], [1, 2]]) self.assertAllEqual(sv[-1:-3:-1], [[15, 16], [13, 14]]) self.assertAllEqual(sv[::-5], [[15, 16], [5, 6]]) self.assertAllEqual(sv[6::-6], [[13, 14], [1, 2]]) self.assertAllEqual(sv[6:5:-6], [[13, 14]]) self.assertAllEqual(sv[6::-7], [[13, 14]]) self.assertAllEqual(sv[7:1:-1], [[15, 16], [13, 14], [11, 12], [9, 10], [7, 8], [5, 6]]) self.assertAllEqual(sv[7:1:-2], [[15, 16], [11, 12], [7, 8]]) self.assertAllEqual(sv[7:1:-4], [[15, 16], [7, 8]]) # Test cases: empty slice self.assertAllEqual(sv[0:0], empty) self.assertAllEqual(sv[5:3], empty) self.assertAllEqual(sv[3:5:-1], empty) self.assertAllEqual(sv[-1:0], empty) self.assertAllEqual(sv[2:-1:-1], empty) # Test cases: slicing other dimensions self.assertAllEqual(sv[:, 0], [1, 3, 5, 7, 9, 11, 13, 15]) self.assertAllEqual(sv[:, 0:1], [[1], [3], [5], [7], [9], [11], [13], [15]]) # Test cases: normal indexing self.assertAllEqual(sv[2], [5, 6]) self.assertAllEqual(sv[6], [13, 14]) self.assertAllEqual(sv[2, 1], 6) self.assertAllEqual(sv[-2], [13, 14]) with self.assertRaisesRegex(IndexError, 'out of bounds'): _ = sv[100] with self.assertRaisesRegex(IndexError, 'out of bounds'): _ = sv[-100] # Test cases: Ellipsis self.assertAllEqual(sv[...], array_ops.concat(v, axis=0)) self.assertAllEqual(sv[..., 0], [1, 3, 5, 7, 9, 11, 13, 15]) self.assertAllEqual(sv[0:1, ...], [[1, 2]]) # Test cases: newaxis self.assertAllEqual( sv[array_ops.newaxis, ...], array_ops.expand_dims_v2(array_ops.concat(v, axis=0), axis=0)) # Test cases: boolean masks self.assertAllEqual(sv[ops.convert_to_tensor(sv) > 10], [11, 12, 13, 14, 15, 16]) # Test cases: tensor input with self.assertRaisesRegex(TypeError, 'not allowed'): _ = sv[constant_op.constant(1)::] with self.assertRaisesRegex(TypeError, 'not allowed'): _ = sv[:constant_op.constant(1):] with self.assertRaisesRegex(TypeError, 'not allowed'): _ = sv[constant_op.constant(1)] # Test cases: inside tf.function @def_function.function def func(): a = sv[:, 0] return a self.assertAllEqual(func(), [1, 3, 5, 7, 9, 11, 13, 15])
def _expand_single_1d_tensor(t): if (hasattr(t, "shape") and isinstance(t.shape, tensor_shape.TensorShape) and t.shape.rank == 1): return array_ops.expand_dims_v2(t, axis=-1) return t