def solve(self, rhs, adjoint=False, adjoint_arg=False, name="solve"):
"""Solve (exact or approx) `R` (batch) systems of equations: `A X = rhs`.
The returned `Tensor` will be close to an exact solution if `A` is well
conditioned. Otherwise closeness will vary. See class docstring for details.
Examples:
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
# Solve R > 0 linear systems for every member of the batch.
RHS = ... # shape [..., M, R]
X = operator.solve(RHS)
# X[..., :, r] is the solution to the r'th linear system
# sum_j A[..., :, j] X[..., j, r] = RHS[..., :, r]
operator.matmul(X)
==> RHS
```
Args:
rhs: `Tensor` with same `dtype` as this operator and compatible shape.
`rhs` is treated like a [batch] matrix meaning for every set of leading
dimensions, the last two dimensions defines a matrix.
See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, solve the system involving the adjoint
of this `LinearOperator`: `A^H X = rhs`.
adjoint_arg: Python `bool`. If `True`, solve `A X = rhs^H` where `rhs^H`
is the hermitian transpose (transposition and complex conjugation).
name: A name scope to use for ops added by this method.
Returns:
`Tensor` with shape `[...,N, R]` and same `dtype` as `rhs`.
Raises:
NotImplementedError: If `self.is_non_singular` or `is_square` is False.
"""
if self.is_non_singular is False:
raise NotImplementedError(
"Exact solve not implemented for an operator that is expected to "
"be singular.")
if self.is_square is False:
raise NotImplementedError(
"Exact solve not implemented for an operator that is expected to "
"not be square.")
with self._name_scope(name, values=[rhs]):
rhs = ops.convert_to_tensor(rhs, name="rhs")
self._check_input_dtype(rhs)
self_dim = -1 if adjoint else -2
arg_dim = -1 if adjoint_arg else -2
tensor_shape.dimension_at_index(
self.shape,
self_dim).assert_is_compatible_with(rhs.get_shape()[arg_dim])
return self._solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
def matvec(self, x, adjoint=False, name="matvec"):
"""Transform [batch] vector `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matric A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
X = ... # shape [..., N], batch vector
Y = operator.matvec(X)
Y.shape
==> [..., M]
Y[..., :] = sum_j A[..., :, j] X[..., j]
```
Args:
x: `Tensor` with compatible shape and same `dtype` as `self`.
`x` is treated as a [batch] vector meaning for every set of leading
dimensions, the last dimension defines a vector.
See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`.
name: A name for this `Op`.
Returns:
A `Tensor` with shape `[..., M]` and same `dtype` as `self`.
"""
with self._name_scope(name):
x = ops.convert_to_tensor(x, name="x")
self._check_input_dtype(x)
self_dim = -2 if adjoint else -1
tensor_shape.dimension_at_index(
self.shape, self_dim).assert_is_compatible_with(x.get_shape()[-1])
return self._matvec(x, adjoint=adjoint)
def solve(self, rhs, adjoint=False, adjoint_arg=False, name="solve"):
"""Solve (exact or approx) `R` (batch) systems of equations: `A X = rhs`.
The returned `Tensor` will be close to an exact solution if `A` is well
conditioned. Otherwise closeness will vary. See class docstring for details.
Examples:
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
# Solve R > 0 linear systems for every member of the batch.
RHS = ... # shape [..., M, R]
X = operator.solve(RHS)
# X[..., :, r] is the solution to the r'th linear system
# sum_j A[..., :, j] X[..., j, r] = RHS[..., :, r]
operator.matmul(X)
==> RHS
```
Args:
rhs: `Tensor` with same `dtype` as this operator and compatible shape.
`rhs` is treated like a [batch] matrix meaning for every set of leading
dimensions, the last two dimensions defines a matrix.
See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, solve the system involving the adjoint
of this `LinearOperator`: `A^H X = rhs`.
adjoint_arg: Python `bool`. If `True`, solve `A X = rhs^H` where `rhs^H`
is the hermitian transpose (transposition and complex conjugation).
name: A name scope to use for ops added by this method.
Returns:
`Tensor` with shape `[...,N, R]` and same `dtype` as `rhs`.
Raises:
NotImplementedError: If `self.is_non_singular` or `is_square` is False.
"""
if self.is_non_singular is False:
raise NotImplementedError(
"Exact solve not implemented for an operator that is expected to "
"be singular.")
if self.is_square is False:
raise NotImplementedError(
"Exact solve not implemented for an operator that is expected to "
"not be square.")
with self._name_scope(name, values=[rhs]):
rhs = ops.convert_to_tensor(rhs, name="rhs")
self._check_input_dtype(rhs)
self_dim = -1 if adjoint else -2
arg_dim = -1 if adjoint_arg else -2
tensor_shape.dimension_at_index(
self.shape, self_dim).assert_is_compatible_with(
rhs.get_shape()[arg_dim])
return self._solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
def static_nrows(self):
"""The number of rows in this partition, if statically known.
```python
self.row_lengths().shape == [self.static_nrows]
self.row_starts().shape == [self.static_nrows]
self.row_limits().shape == [self.static_nrows]
self.row_splits().shape == [self.static_nrows + 1]
```
Returns:
The number of rows in this partition as an `int` (if statically known);
or `None` (otherwise).
"""
if self._row_splits is not None:
nrows = tensor_shape.dimension_at_index(self._row_splits.shape,
0) - 1
if nrows.value is not None:
return nrows
if self._row_lengths is not None:
nrows = tensor_shape.dimension_at_index(self._row_lengths.shape, 0)
if nrows.value is not None:
return nrows
if self._nrows is not None:
return tensor_shape.Dimension(
tensor_util.constant_value(self._nrows))
return None
def matvec(self, x, adjoint=False, name="matvec"):
"""Transform [batch] vector `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matric A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
X = ... # shape [..., N], batch vector
Y = operator.matvec(X)
Y.shape
==> [..., M]
Y[..., :] = sum_j A[..., :, j] X[..., j]
```
Args:
x: `Tensor` with compatible shape and same `dtype` as `self`.
`x` is treated as a [batch] vector meaning for every set of leading
dimensions, the last dimension defines a vector.
See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`.
name: A name for this `Op`.
Returns:
A `Tensor` with shape `[..., M]` and same `dtype` as `self`.
"""
with self._name_scope(name):
x = ops.convert_to_tensor(x, name="x")
self._check_input_dtype(x)
self_dim = -2 if adjoint else -1
tensor_shape.dimension_at_index(
self.shape,
self_dim).assert_is_compatible_with(x.get_shape()[-1])
return self._matvec(x, adjoint=adjoint)
def _verify_input(tensor_list, labels, probs_list):
"""Verify that batched inputs are well-formed."""
checked_probs_list = []
for probs in probs_list:
# Since number of classes shouldn't change at runtime, probabilities shape
# should be fully defined.
probs.get_shape().assert_is_fully_defined()
# Probabilities must be 1D.
probs.get_shape().assert_has_rank(1)
# Probabilities must be nonnegative and sum to one.
tol = 1e-6
prob_sum = math_ops.reduce_sum(probs)
checked_probs = control_flow_ops.with_dependencies([
check_ops.assert_non_negative(probs),
check_ops.assert_less(prob_sum, 1.0 + tol),
check_ops.assert_less(1.0 - tol, prob_sum)
], probs)
checked_probs_list.append(checked_probs)
# All probabilities should be the same length.
prob_length = checked_probs_list[0].get_shape().num_elements()
for checked_prob in checked_probs_list:
if checked_prob.get_shape().num_elements() != prob_length:
raise ValueError('Probability parameters must have the same length.')
# Labels tensor should only have batch dimension.
labels.get_shape().assert_has_rank(1)
for tensor in tensor_list:
# Data tensor should have a batch dimension.
shape = tensor.get_shape().with_rank_at_least(1)
# Data and label batch dimensions must be compatible.
tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with(
labels.get_shape()[0])
# Data and labels must have the same, strictly positive batch size. Since we
# can't assume we know the batch size at graph creation, add runtime checks.
labels_batch_size = array_ops.shape(labels)[0]
lbl_assert = check_ops.assert_positive(labels_batch_size)
# Make each tensor depend on its own checks.
labels = control_flow_ops.with_dependencies([lbl_assert], labels)
tensor_list = [
control_flow_ops.with_dependencies([
lbl_assert,
check_ops.assert_equal(array_ops.shape(x)[0], labels_batch_size)
], x) for x in tensor_list
]
# Label's classes must be integers 0 <= x < num_classes.
labels = control_flow_ops.with_dependencies([
check_ops.assert_integer(labels), check_ops.assert_non_negative(labels),
check_ops.assert_less(labels, math_ops.cast(prob_length, labels.dtype))
], labels)
return tensor_list, labels, checked_probs_list
def _verify_input(tensor_list, labels, probs_list):
"""Verify that batched inputs are well-formed."""
checked_probs_list = []
for probs in probs_list:
# Since number of classes shouldn't change at runtime, probabilities shape
# should be fully defined.
probs.get_shape().assert_is_fully_defined()
# Probabilities must be 1D.
probs.get_shape().assert_has_rank(1)
# Probabilities must be nonnegative and sum to one.
tol = 1e-6
prob_sum = math_ops.reduce_sum(probs)
checked_probs = control_flow_ops.with_dependencies([
check_ops.assert_non_negative(probs),
check_ops.assert_less(prob_sum, 1.0 + tol),
check_ops.assert_less(1.0 - tol, prob_sum)
], probs)
checked_probs_list.append(checked_probs)
# All probabilities should be the same length.
prob_length = checked_probs_list[0].get_shape().num_elements()
for checked_prob in checked_probs_list:
if checked_prob.get_shape().num_elements() != prob_length:
raise ValueError('Probability parameters must have the same length.')
# Labels tensor should only have batch dimension.
labels.get_shape().assert_has_rank(1)
for tensor in tensor_list:
# Data tensor should have a batch dimension.
shape = tensor.get_shape().with_rank_at_least(1)
# Data and label batch dimensions must be compatible.
tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with(
labels.get_shape()[0])
# Data and labels must have the same, strictly positive batch size. Since we
# can't assume we know the batch size at graph creation, add runtime checks.
labels_batch_size = array_ops.shape(labels)[0]
lbl_assert = check_ops.assert_positive(labels_batch_size)
# Make each tensor depend on its own checks.
labels = control_flow_ops.with_dependencies([lbl_assert], labels)
tensor_list = [
control_flow_ops.with_dependencies([
lbl_assert,
check_ops.assert_equal(array_ops.shape(x)[0], labels_batch_size)
], x) for x in tensor_list
]
# Label's classes must be integers 0 <= x < num_classes.
labels = control_flow_ops.with_dependencies([
check_ops.assert_integer(labels), check_ops.assert_non_negative(labels),
check_ops.assert_less(labels, math_ops.cast(prob_length, labels.dtype))
], labels)
return tensor_list, labels, checked_probs_list
def matmul(self, x, adjoint=False, adjoint_arg=False, name="matmul"):
"""Transform [batch] matrix `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
X = ... # shape [..., N, R], batch matrix, R > 0.
Y = operator.matmul(X)
Y.shape
==> [..., M, R]
Y[..., :, r] = sum_j A[..., :, j] X[j, r]
```
Args:
x: `LinearOperator` or `Tensor` with compatible shape and same `dtype` as
`self`. See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`.
adjoint_arg: Python `bool`. If `True`, compute `A x^H` where `x^H` is
the hermitian transpose (transposition and complex conjugation).
name: A name for this `Op`.
Returns:
A `LinearOperator` or `Tensor` with shape `[..., M, R]` and same `dtype`
as `self`.
"""
if isinstance(x, LinearOperator):
left_operator = self.adjoint() if adjoint else self
right_operator = x.adjoint() if adjoint_arg else x
if (right_operator.range_dimension is not None
and left_operator.domain_dimension is not None
and right_operator.range_dimension !=
left_operator.domain_dimension):
raise ValueError(
"Operators are incompatible. Expected `x` to have dimension"
" {} but got {}.".format(left_operator.domain_dimension,
right_operator.range_dimension))
with self._name_scope(name):
return linear_operator_algebra.matmul(left_operator,
right_operator)
with self._name_scope(name):
x = ops.convert_to_tensor(x, name="x")
self._check_input_dtype(x)
self_dim = -2 if adjoint else -1
arg_dim = -1 if adjoint_arg else -2
tensor_shape.dimension_at_index(
self.shape,
self_dim).assert_is_compatible_with(x.get_shape()[arg_dim])
return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _verify_data_inputs(tensor_list):
"""Verify that batched data inputs are well-formed."""
for tensor in tensor_list:
# Data tensor should have a batch dimension.
shape = tensor.get_shape().with_rank_at_least(1)
# Data batch dimensions must be compatible.
tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with(
tensor_list[0].get_shape()[0])
return tensor_list
def matmul(self, x, adjoint=False, adjoint_arg=False, name="matmul"):
"""Transform [batch] matrix `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
X = ... # shape [..., N, R], batch matrix, R > 0.
Y = operator.matmul(X)
Y.shape
==> [..., M, R]
Y[..., :, r] = sum_j A[..., :, j] X[j, r]
```
Args:
x: `LinearOperator` or `Tensor` with compatible shape and same `dtype` as
`self`. See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`.
adjoint_arg: Python `bool`. If `True`, compute `A x^H` where `x^H` is
the hermitian transpose (transposition and complex conjugation).
name: A name for this `Op`.
Returns:
A `LinearOperator` or `Tensor` with shape `[..., M, R]` and same `dtype`
as `self`.
"""
if isinstance(x, LinearOperator):
left_operator = self.adjoint() if adjoint else self
right_operator = x.adjoint() if adjoint_arg else x
if (right_operator.range_dimension is not None and
left_operator.domain_dimension is not None and
right_operator.range_dimension != left_operator.domain_dimension):
raise ValueError(
"Operators are incompatible. Expected `x` to have dimension"
" {} but got {}.".format(
left_operator.domain_dimension, right_operator.range_dimension))
with self._name_scope(name):
return linear_operator_algebra.matmul(left_operator, right_operator)
with self._name_scope(name):
x = ops.convert_to_tensor(x, name="x")
self._check_input_dtype(x)
self_dim = -2 if adjoint else -1
arg_dim = -1 if adjoint_arg else -2
tensor_shape.dimension_at_index(
self.shape, self_dim).assert_is_compatible_with(
x.get_shape()[arg_dim])
return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _verify_data_inputs(tensor_list):
"""Verify that batched data inputs are well-formed."""
for tensor in tensor_list:
# Data tensor should have a batch dimension.
shape = tensor.get_shape().with_rank_at_least(1)
# Data batch dimensions must be compatible.
tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with(
tensor_list[0].get_shape()[0])
return tensor_list
def solvevec(self, rhs, adjoint=False, name="solve"):
"""Solve single equation with best effort: `A X = rhs`.
The returned `Tensor` will be close to an exact solution if `A` is well
conditioned. Otherwise closeness will vary. See class docstring for details.
Examples:
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
# Solve one linear system for every member of the batch.
RHS = ... # shape [..., M]
X = operator.solvevec(RHS)
# X is the solution to the linear system
# sum_j A[..., :, j] X[..., j] = RHS[..., :]
operator.matvec(X)
==> RHS
```
Args:
rhs: `Tensor` with same `dtype` as this operator.
`rhs` is treated like a [batch] vector meaning for every set of leading
dimensions, the last dimension defines a vector. See class docstring
for definition of compatibility regarding batch dimensions.
adjoint: Python `bool`. If `True`, solve the system involving the adjoint
of this `LinearOperator`: `A^H X = rhs`.
name: A name scope to use for ops added by this method.
Returns:
`Tensor` with shape `[...,N]` and same `dtype` as `rhs`.
Raises:
NotImplementedError: If `self.is_non_singular` or `is_square` is False.
"""
with self._name_scope(name):
rhs = ops.convert_to_tensor(rhs, name="rhs")
self._check_input_dtype(rhs)
self_dim = -1 if adjoint else -2
tensor_shape.dimension_at_index(
self.shape, self_dim).assert_is_compatible_with(
rhs.get_shape()[-1])
return self._solvevec(rhs, adjoint=adjoint)
def solvevec(self, rhs, adjoint=False, name="solve"):
"""Solve single equation with best effort: `A X = rhs`.
The returned `Tensor` will be close to an exact solution if `A` is well
conditioned. Otherwise closeness will vary. See class docstring for details.
Examples:
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
# Solve one linear system for every member of the batch.
RHS = ... # shape [..., M]
X = operator.solvevec(RHS)
# X is the solution to the linear system
# sum_j A[..., :, j] X[..., j] = RHS[..., :]
operator.matvec(X)
==> RHS
```
Args:
rhs: `Tensor` with same `dtype` as this operator.
`rhs` is treated like a [batch] vector meaning for every set of leading
dimensions, the last dimension defines a vector. See class docstring
for definition of compatibility regarding batch dimensions.
adjoint: Python `bool`. If `True`, solve the system involving the adjoint
of this `LinearOperator`: `A^H X = rhs`.
name: A name scope to use for ops added by this method.
Returns:
`Tensor` with shape `[...,N]` and same `dtype` as `rhs`.
Raises:
NotImplementedError: If `self.is_non_singular` or `is_square` is False.
"""
with self._name_scope(name):
rhs = ops.convert_to_tensor(rhs, name="rhs")
self._check_input_dtype(rhs)
self_dim = -1 if adjoint else -2
tensor_shape.dimension_at_index(
self.shape,
self_dim).assert_is_compatible_with(rhs.get_shape()[-1])
return self._solvevec(rhs, adjoint=adjoint)
def _merge_nrows(nrows, static_nrows, value, dtype, validate):
"""Merges `nrows` with `nrows(value)`.
Checks that `value` has the expected number of rows (`nrows`), and returns
`nrows`. If `validate` is true, then add validation ops that check that
the `nrows` values match.
Args:
nrows: scalar integer Tensor.
static_nrows: tf.Dimension: static value of nrows, if known.
value: Tensor or RaggedTensor or StructuredTensor
dtype: dtype for `nrows`.
validate: bool -- whether to add validation ops.
Returns:
A tuple `(nrows, static_nrows)`.
"""
static_value_nrows = tensor_shape.dimension_at_index(value.shape, 0)
if isinstance(value, ops.Tensor):
value_nrows = array_ops.shape(value, out_type=dtype)[0]
else:
value_nrows = value.nrows()
if nrows is None:
nrows = value_nrows
elif (static_value_nrows.value is not None
and static_nrows.value is not None):
if not static_value_nrows.is_compatible_with(static_nrows):
raise ValueError('fields have incompatible nrows')
nrows = value_nrows # No need to add an assertion op.
elif validate:
nrows = control_flow_ops.with_dependencies([
check_ops.assert_equal(
nrows, value_nrows, message='fields have incompatible nrows')
], nrows)
return nrows, static_nrows.merge_with(static_value_nrows)
def nrows(self, out_type=None, name=None):
"""Returns the number of rows in this ragged tensor.
I.e., the size of the outermost dimension of the tensor.
Args:
out_type: `dtype` for the returned tensor. Defaults to
`self.row_splits.dtype`.
name: A name prefix for the returned tensor (optional).
Returns:
A scalar `Tensor` with dtype `out_type`.
"""
if out_type is None:
out_type = self._row_splits.dtype
else:
out_type = dtypes.as_dtype(out_type)
if self._cached_nrows is not None:
return math_ops.cast(self._cached_nrows, out_type)
with ops.name_scope(name, "RaggedNRows", [self]):
nsplits = tensor_shape.dimension_at_index(self.row_splits.shape, 0)
if nsplits.value is None:
return array_ops.shape(self.row_splits,
out_type=out_type)[0] - 1
else:
return constant_op.constant(nsplits.value - 1, dtype=out_type)
def _check_shapes(self):
"""Static check that shapes are compatible."""
# Broadcast shape also checks that u and v are compatible.
uv_shape = array_ops.broadcast_static_shape(
self.u.get_shape(), self.v.get_shape())
batch_shape = array_ops.broadcast_static_shape(
self.base_operator.batch_shape, uv_shape[:-2])
tensor_shape.Dimension(
self.base_operator.domain_dimension).assert_is_compatible_with(
uv_shape[-2])
if self._diag_update is not None:
tensor_shape.dimension_at_index(uv_shape, -1).assert_is_compatible_with(
self._diag_update.get_shape()[-1])
array_ops.broadcast_static_shape(
batch_shape, self._diag_update.get_shape()[:-1])
def _check_shapes(self):
"""Static check that shapes are compatible."""
# Broadcast shape also checks that u and v are compatible.
uv_shape = array_ops.broadcast_static_shape(self.u.shape, self.v.shape)
batch_shape = array_ops.broadcast_static_shape(
self.base_operator.batch_shape, uv_shape[:-2])
tensor_shape.Dimension(
self.base_operator.domain_dimension).assert_is_compatible_with(
uv_shape[-2])
if self._diag_update is not None:
tensor_shape.dimension_at_index(uv_shape,
-1).assert_is_compatible_with(
self._diag_update.shape[-1])
array_ops.broadcast_static_shape(batch_shape,
self._diag_update.shape[:-1])
def ragged_tensor_spec(shape=None,
dtype=dtypes.float32,
ragged_rank=None,
row_splits_dtype=dtypes.int64,
name=None):
"""Returns a tensor specification for a RaggedTensor.
Returns an object which can be passed to `tf.function` (or other
functions that expect `TensorSpec`s) to specify shape constraints
for a `RaggedTensor` argument.
Args:
shape: The shape of the RaggedTensor, or `None` to allow any shape.
dtype: Data type of values in the RaggedTensor.
ragged_rank: Python integer, the ragged rank of the RaggedTensor
to be described. Defaults to `shape.ndims - 1`.
row_splits_dtype: `dtype` for the RaggedTensor's `row_splits` tensor.
One of `tf.int32` or `tf.int64`.
name: Optional name prefix for the `TensorSpec`s.
Returns:
An object describing the `flat_values` and `nested_row_splits` tensors
that comprise the `RaggedTensor`.
"""
dtype = dtypes.as_dtype(dtype)
shape = tensor_shape.TensorShape(shape)
if ragged_rank is None:
if shape.ndims is None:
raise ValueError(
"Must specify ragged_rank or a shape with known rank.")
ragged_rank = shape.ndims - 1
elif not isinstance(ragged_rank, int):
raise TypeError("ragged_rank must be an int")
if ragged_rank == 0:
return tensor_spec.TensorSpec(shape=shape, dtype=dtype, name=name)
result = tensor_spec.TensorSpec(
tensor_shape.TensorShape([None]).concatenate(shape[ragged_rank + 1:]),
dtype, name)
for i in range(ragged_rank - 1, 0, -1):
splits = tensor_spec.TensorSpec([None], row_splits_dtype,
"%s.row_splits_%d" %
(name, i) if name else None)
result = ragged_tensor.RaggedTensor.from_row_splits(result, splits)
outer_dim = tensor_shape.dimension_at_index(shape, 0)
splits_shape = [None if outer_dim is None else outer_dim + 1]
splits = tensor_spec.TensorSpec(splits_shape, row_splits_dtype,
"%s.row_splits_0" % name if name else None)
result = ragged_tensor.RaggedTensor.from_row_splits(result, splits)
return result
def matmul(self, x, adjoint=False, adjoint_arg=False, name="matmul"):
"""Transform [batch] matrix `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
X = ... # shape [..., N, R], batch matrix, R > 0.
Y = operator.matmul(X)
Y.shape
==> [..., M, R]
Y[..., :, r] = sum_j A[..., :, j] X[j, r]
```
Args:
x: `Tensor` with compatible shape and same `dtype` as `self`.
See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`.
adjoint_arg: Python `bool`. If `True`, compute `A x^H` where `x^H` is
the hermitian transpose (transposition and complex conjugation).
name: A name for this `Op.
Returns:
A `Tensor` with shape `[..., M, R]` and same `dtype` as `self`.
"""
with self._name_scope(name, values=[x]):
x = ops.convert_to_tensor(x, name="x")
self._check_input_dtype(x)
self_dim = -2 if adjoint else -1
arg_dim = -1 if adjoint_arg else -2
tensor_shape.dimension_at_index(
self.shape,
self_dim).assert_is_compatible_with(x.get_shape()[arg_dim])
return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def shape(self):
"""The statically known shape of this RaggedStructuredTensor."""
nrows = tensor_shape.dimension_at_index(self._row_splits.shape, 0) - 1
if self._uniform_row_length is not None:
row_length = tensor_util.constant_value(self._uniform_row_length)
else:
row_length = None
values_shape = self._values.shape
value_shape = values_shape[1:]
return tensor_shape.TensorShape([nrows,
row_length]).concatenate(value_shape)
def matmul(self, x, adjoint=False, adjoint_arg=False, name="matmul"):
"""Transform [batch] matrix `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matrix A. Assume A.shape = [..., M, N]
operator = LinearOperator(...)
operator.shape = [..., M, N]
X = ... # shape [..., N, R], batch matrix, R > 0.
Y = operator.matmul(X)
Y.shape
==> [..., M, R]
Y[..., :, r] = sum_j A[..., :, j] X[j, r]
```
Args:
x: `Tensor` with compatible shape and same `dtype` as `self`.
See class docstring for definition of compatibility.
adjoint: Python `bool`. If `True`, left multiply by the adjoint: `A^H x`.
adjoint_arg: Python `bool`. If `True`, compute `A x^H` where `x^H` is
the hermitian transpose (transposition and complex conjugation).
name: A name for this `Op.
Returns:
A `Tensor` with shape `[..., M, R]` and same `dtype` as `self`.
"""
with self._name_scope(name, values=[x]):
x = ops.convert_to_tensor(x, name="x")
self._check_input_dtype(x)
self_dim = -2 if adjoint else -1
arg_dim = -1 if adjoint_arg else -2
tensor_shape.dimension_at_index(
self.shape, self_dim).assert_is_compatible_with(
x.get_shape()[arg_dim])
return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def ragged_tensor_spec(shape=None, dtype=dtypes.float32,
ragged_rank=None, row_splits_dtype=dtypes.int64,
name=None):
"""Returns a tensor specification for a RaggedTensor.
Returns an object which can be passed to `tf.function` (or other
functions that expect `TensorSpec`s) to specify shape constraints
for a `RaggedTensor` argument.
Args:
shape: The shape of the RaggedTensor, or `None` to allow any shape.
dtype: Data type of values in the RaggedTensor.
ragged_rank: Python integer, the ragged rank of the RaggedTensor
to be described. Defaults to `shape.ndims - 1`.
row_splits_dtype: `dtype` for the RaggedTensor's `row_splits` tensor.
One of `tf.int32` or `tf.int64`.
name: Optional name prefix for the `TensorSpec`s.
Returns:
An object describing the `flat_values` and `nested_row_splits` tensors
that comprise the `RaggedTensor`.
"""
dtype = dtypes.as_dtype(dtype)
shape = tensor_shape.TensorShape(shape)
if ragged_rank is None:
if shape.ndims is None:
raise ValueError("Must specify ragged_rank or a shape with known rank.")
ragged_rank = shape.ndims - 1
elif not isinstance(ragged_rank, int):
raise TypeError("ragged_rank must be an int")
if ragged_rank == 0:
return tensor_spec.TensorSpec(shape=shape, dtype=dtype, name=name)
result = tensor_spec.TensorSpec(
tensor_shape.TensorShape([None]).concatenate(shape[ragged_rank + 1:]),
dtype, name)
for i in range(ragged_rank - 1, 0, -1):
splits = tensor_spec.TensorSpec(
[None], row_splits_dtype,
"%s.row_splits_%d" % (name, i) if name else None)
result = ragged_tensor.RaggedTensor.from_row_splits(result, splits)
outer_dim = tensor_shape.dimension_at_index(shape, 0)
splits_shape = [None if outer_dim is None else outer_dim + 1]
splits = tensor_spec.TensorSpec(
splits_shape, row_splits_dtype,
"%s.row_splits_0" % name if name else None)
result = ragged_tensor.RaggedTensor.from_row_splits(result, splits)
return result
def static_nvals(self):
"""The number of values in this partition, if statically known.
```python
self.value_rowids().shape == [self.static_vals]
```
Returns:
The number of values in this partition as an `int` (if statically known);
or `None` (otherwise).
"""
if self._value_rowids is not None:
nvals = tensor_shape.dimension_at_index(self._value_rowids.shape, 0)
if nvals.value is not None:
return nvals.value
return None
def _replace_ragged_with_flat_values(value, partition_lists,
flat_values_nrows):
"""Replace RaggedTensors with their flat_values, and record their partitions.
Returns a copy of `value`, with any nested `RaggedTensor`s replaced by their
`flat_values` tensor. Looks inside lists, tuples, and dicts.
Appends each `RaggedTensor`'s `RowPartition`s to `partition_lists`.
Args:
value: The value that should be transformed by replacing `RaggedTensors`.
partition_lists: An output parameter used to record the row partitions
for any `RaggedTensors` that were replaced.
flat_values_nrows: An output parameter used to record the outer dimension
size for each replacement `flat_values` (when known). Contains a list of
int.
Returns:
A copy of `value` with nested `RaggedTensors` replaced by their `values`.
"""
# Base case
if ragged_tensor.is_ragged(value):
value = ragged_tensor.convert_to_tensor_or_ragged_tensor(value)
partition_lists.append(value._nested_row_partitions) # pylint: disable=protected-access
nrows = tensor_shape.dimension_at_index(value.flat_values.shape,
0).value
if nrows is not None:
flat_values_nrows.append(nrows)
return value.flat_values
# Recursion cases
def recurse(v):
return _replace_ragged_with_flat_values(v, partition_lists,
flat_values_nrows)
if isinstance(value, list):
return [recurse(v) for v in value]
elif isinstance(value, tuple):
return tuple(recurse(v) for v in value)
elif isinstance(value, dict):
return dict((k, recurse(v)) for (k, v) in value.items())
else:
return value
def nrows(self, out_type=None):
"""Returns the number of rows created by this `RowPartition`.
Args:
out_type: `dtype` for the returned tensor. Defaults to `self.dtype`.
Returns:
scalar integer Tensor
"""
if out_type is None:
out_type = self.dtype
else:
out_type = dtypes.as_dtype(out_type)
if self._nrows is not None:
return math_ops.cast(self._nrows, out_type)
nsplits = tensor_shape.dimension_at_index(self._row_splits.shape, 0)
if nsplits.value is None:
return array_ops.shape(self._row_splits, out_type=out_type)[0] - 1
else:
return constant_op.constant(nsplits.value - 1, dtype=out_type)
def shape(self):
"""The statically known shape of this ragged tensor.
Returns:
A `TensorShape` containing the statically known shape of this ragged
tensor. Ragged dimensions have a size of `None`.
Examples:
```python
>>> ragged.constant([[0], [1, 2]]).shape
TensorShape([Dimension(2), Dimension(None)])
>>> ragged.constant([[[0, 1]], [[1, 2], [3, 4]]], ragged_rank=1).shape
TensorShape([Dimension(2), Dimension(None), Dimension(2)
```
"""
nrows = tensor_shape.dimension_at_index(self._row_splits.shape, 0) - 1
values_shape = self._values.shape
value_shape = values_shape[1:]
return tensor_shape.TensorShape([nrows, None]).concatenate(value_shape)
def shape(self):
"""The statically known shape of this ragged tensor.
Returns:
A `TensorShape` containing the statically known shape of this ragged
tensor. Ragged dimensions have a size of `None`.
Examples:
```python
>>> ragged.constant([[0], [1, 2]]).shape
TensorShape([Dimension(2), Dimension(None)])
>>> ragged.constant([[[0, 1]], [[1, 2], [3, 4]]], ragged_rank=1).shape
TensorShape([Dimension(2), Dimension(None), Dimension(2)
```
"""
nrows = tensor_shape.dimension_at_index(self._row_splits.shape, 0) - 1
values_shape = self._values.shape
value_shape = values_shape[1:]
return tensor_shape.TensorShape([nrows, None]).concatenate(value_shape)
def nrows(self, out_type=None, name=None):
"""Returns the number of rows created by this `RowPartition`.
Args:
out_type: `dtype` for the returned tensor. Defaults to
`self.dtype`.
name: A name prefix for the returned tensor (optional).
Returns:
scalar integer Tensor
"""
if out_type is None:
out_type = self.dtype
else:
out_type = dtypes.as_dtype(out_type)
if self._cached_nrows is not None:
return math_ops.cast(self._cached_nrows, out_type)
with ops.name_scope(name, "RaggedNRows", [self]):
nsplits = tensor_shape.dimension_at_index(self.row_splits.shape, 0)
if nsplits.value is None:
return array_ops.shape(self.row_splits, out_type=out_type)[0] - 1
else:
return constant_op.constant(nsplits.value - 1, dtype=out_type)
def _ragged_getitem(rt_input, key_list):
"""Helper for indexing and slicing ragged tensors with __getitem__().
Extracts the specified piece of the `rt_input`. See
`RaggedTensor.__getitem__` for examples and restrictions.
Args:
rt_input: The `RaggedTensor` from which a piece should be returned.
key_list: The list of keys specifying which piece to return. Each key
corresponds with a separate dimension.
Returns:
The indicated piece of rt_input.
Raises:
ValueError: If `key_list` is not supported.
TypeError: If any keys in `key_list` have an unsupported type.
"""
if not key_list:
return rt_input
row_key = key_list[0]
inner_keys = key_list[1:]
if row_key is Ellipsis:
expanded_key_list = _expand_ellipsis(key_list, rt_input.shape.ndims)
return _ragged_getitem(rt_input, expanded_key_list)
# Adding a new axis: Get rt_input[inner_keys], and wrap it in a RaggedTensor
# that puts all values in a single row.
if row_key is array_ops.newaxis:
inner_rt = _ragged_getitem(rt_input, inner_keys)
nsplits = tensor_shape.dimension_at_index(inner_rt.row_splits.shape, 0)
if nsplits.value is not None:
nsplits = nsplits.value
else:
nsplits = array_ops.shape(inner_rt.row_splits,
out_type=inner_rt.row_splits.dtype)[0]
return ragged_tensor.RaggedTensor.from_uniform_row_length(
inner_rt, nsplits - 1, nrows=1, validate=False)
# Slicing a range of rows: first slice the outer dimension, and then
# call `_ragged_getitem_inner_dimensions` to handle the inner keys.
if isinstance(row_key, slice):
sliced_rt_input = _slice_ragged_row_dimension(rt_input, row_key)
if rt_input.uniform_row_length is not None:
# If the inner dimension has uniform_row_length, then preserve it (by
# re-wrapping the values in a new RaggedTensor). Note that the row
# length won't have changed, since we're slicing a range of rows (and not
# slicing the rows themselves).
sliced_rt_input = ragged_tensor.RaggedTensor.from_uniform_row_length(
sliced_rt_input.values,
rt_input.uniform_row_length,
nrows=sliced_rt_input.nrows())
return _ragged_getitem_inner_dimensions(sliced_rt_input, inner_keys)
# Indexing a single row: slice values to get the indicated row, and then
# use a recursive call to __getitem__ to handle the inner keys.
else:
starts = rt_input.row_splits[:-1]
limits = rt_input.row_splits[1:]
if context.executing_eagerly():
# In python, __getitem__ should throw IndexError for out of bound
# indices. This will allow iteration run correctly as python will
# translate IndexError into StopIteration for next()/__next__().
# Below is an example:
# import tensorflow as tf
# r = tf.ragged.constant([[1., 2.], [3., 4., 5.], [6.]])
# for elem in r:
# print(elem)
# In non eager mode, the exception is thrown when session runs
# so we don't know if out of bound happens before.
# In eager mode, however, it is possible to find out when to
# throw out of bound IndexError.
# In the following row_key >= len(starts) is checked. In case of
# TypeError which happens when row_key is not an integer, the exception
# will simply be ignored as it will be processed later anyway.
try:
if int(row_key) >= len(starts):
raise IndexError(
"Row key {} out of bounds".format(row_key))
except (TypeError, ValueError):
pass
row = rt_input.values[starts[row_key]:limits[row_key]]
return row.__getitem__(inner_keys)
def nrows_as_dimension(self): """Returns the first dimension of the shape as a `tf.Dimension`.""" return tensor_shape.dimension_at_index(self._row_splits.shape, 0) - 1
def _component_specs(self):
row_splits_shape = tensor_shape.TensorShape(
[tensor_shape.dimension_at_index(self._nrows, 0) + 1])
return tensor_spec.TensorSpec(row_splits_shape, self._dtype)
def _ragged_getitem_inner_dimensions(rt_input, key_list):
"""Retrieve inner dimensions, keeping outermost dimension unchanged.
Args:
rt_input: The `RaggedTensor` or `Tensor` from which a piece should be
extracted.
key_list: The __getitem__ keys for slicing the inner dimensions.
Returns:
A `RaggedTensor`.
Raises:
ValueError: If key_list is not supported.
"""
if not key_list:
return rt_input
if isinstance(rt_input, ops.Tensor):
return rt_input.__getitem__([slice(None, None, None)] + key_list)
column_key = key_list[0]
if column_key is Ellipsis:
expanded_key_list = _expand_ellipsis(key_list,
rt_input.values.shape.ndims)
return _ragged_getitem_inner_dimensions(rt_input, expanded_key_list)
# Adding a new axis to a ragged inner dimension: recursively get the inner
# dimensions of rt_input with key_list[1:], and then wrap the result in a
# RaggedTensor that puts each value in its own row.
if column_key is array_ops.newaxis:
inner_rt = _ragged_getitem_inner_dimensions(rt_input, key_list[1:])
nsplits = tensor_shape.dimension_at_index(inner_rt.row_splits.shape, 0)
if nsplits.value is not None:
nsplits = nsplits.value
else:
nsplits = array_ops.shape(inner_rt.row_splits,
out_type=inner_rt.row_splits.dtype)[0]
return ragged_tensor.RaggedTensor.from_uniform_row_length(
inner_rt, 1, nrows=nsplits - 1, validate=False)
# Slicing a range of columns in a ragged inner dimension. We use a
# recursive call to process the values, and then assemble a RaggedTensor
# with those values.
if isinstance(column_key, slice):
if (column_key.start is None and column_key.stop is None
and column_key.step is None):
# Trivial slice: recursively process all values, & splits is unchanged.
return rt_input.with_values(
_ragged_getitem_inner_dimensions(rt_input.values,
key_list[1:]))
else:
if not (isinstance(column_key.start, (ops.Tensor, int, type(None)))
and isinstance(column_key.stop,
(ops.Tensor, int, type(None)))):
raise TypeError("slice offsets must be integers or None")
# Nontrivial slice: use ragged_gather to extract the indicated slice as
# a new RaggedTensor (inner_rt), and then recursively process its values.
starts = rt_input.row_splits[:-1]
limits = rt_input.row_splits[1:]
step = 1 if column_key.step is None else column_key.step
lower_bound = _if_ge_zero(step, lambda: starts, lambda: starts - 1)
upper_bound = _if_ge_zero(step, lambda: limits, lambda: limits - 1)
# inner_rt_starts[i] = index to start gathering for row i.
if column_key.start is None:
inner_rt_starts = _if_ge_zero(step, lambda: starts,
lambda: limits - 1)
else:
start_offset = math_ops.cast(column_key.start, starts.dtype)
inner_rt_starts = _if_ge_zero(
column_key.start, lambda: math_ops.minimum(
starts + start_offset, upper_bound), lambda: math_ops.
maximum(limits + start_offset, lower_bound))
# inner_rt_limits[i] = index to stop gathering for row i.
if column_key.stop is None:
inner_rt_limits = _if_ge_zero(step, lambda: limits,
lambda: starts - 1)
else:
stop_offset = math_ops.cast(column_key.stop, starts.dtype)
inner_rt_limits = _if_ge_zero(
column_key.stop, lambda: math_ops.minimum(
starts + stop_offset, upper_bound), lambda: math_ops.
maximum(limits + stop_offset, lower_bound))
inner_rt = _build_ragged_tensor_from_value_ranges(
inner_rt_starts, inner_rt_limits, column_key.step,
rt_input.values)
# If the row dimension is uniform, then calculate the new
# uniform_row_length, and rebuild inner_rt using that uniform_row_lengths.
if rt_input.uniform_row_length is not None:
new_row_length = _slice_length(rt_input.uniform_row_length,
column_key)
inner_rt = ragged_tensor.RaggedTensor.from_uniform_row_length(
inner_rt.values, new_row_length, rt_input.nrows())
return inner_rt.with_values(
_ragged_getitem_inner_dimensions(inner_rt.values,
key_list[1:]))
# Indexing a single column in a ragged inner dimension: raise an Exception.
# See RaggedTensor.__getitem__.__doc__ for an explanation of why indexing
# into a ragged inner dimension is problematic.
if rt_input.uniform_row_length is None:
raise ValueError("Cannot index into an inner ragged dimension.")
# Indexing a single column in a uniform inner dimension: check that the
# given index is in-bounds, and then use a strided slice over rt_input.values
# to take the indicated element from each row.
row_length = rt_input.uniform_row_length
column_key = math_ops.cast(column_key, row_length.dtype)
oob_err_msg = "Index out of bounds when indexing into a ragged tensor"
oob_checks = [
check_ops.assert_greater_equal(column_key,
-row_length,
message=oob_err_msg),
check_ops.assert_less(column_key, row_length, message=oob_err_msg),
]
with ops.control_dependencies(oob_checks):
offset = _if_ge_zero(column_key, lambda: column_key,
lambda: row_length + column_key)
sliced_rt = rt_input.values[offset::row_length]
return _ragged_getitem_inner_dimensions(sliced_rt, key_list[1:])
def __init__(
self,
input_shape,
filter_shape,
padding,
strides=None,
dilation_rate=None,
name=None,
data_format=None,
):
"""Helper function for convolution."""
num_total_dims = filter_shape.ndims
if num_total_dims is None:
num_total_dims = input_shape.ndims
if num_total_dims is None:
raise ValueError("rank of input or filter must be known")
num_spatial_dims = num_total_dims - 2
try:
input_shape.with_rank(num_spatial_dims + 2)
except ValueError:
raise ValueError("input tensor must have rank %d" %
(num_spatial_dims + 2))
try:
filter_shape.with_rank(num_spatial_dims + 2)
except ValueError:
raise ValueError("filter tensor must have rank %d" %
(num_spatial_dims + 2))
if data_format is None or not data_format.startswith("NC"):
input_channels_dim = tensor_shape.dimension_at_index(
input_shape, num_spatial_dims + 1)
spatial_dims = range(1, num_spatial_dims + 1)
else:
input_channels_dim = tensor_shape.dimension_at_index(
input_shape, 1)
spatial_dims = range(2, num_spatial_dims + 2)
filter_dim = tensor_shape.dimension_at_index(filter_shape,
num_spatial_dims)
if not (input_channels_dim % filter_dim).is_compatible_with(0):
raise ValueError(
"number of input channels is not divisible by corresponding "
"dimension of filter, {} % {} != 0".format(
input_channels_dim, filter_dim))
strides, dilation_rate = nn_ops._get_strides_and_dilation_rate(
num_spatial_dims, strides, dilation_rate)
self.input_shape = input_shape
self.filter_shape = filter_shape
self.data_format = data_format
self.strides = strides
self.padding = padding
self.name = name
self.dilation_rate = dilation_rate
self.conv_op = nn_ops._WithSpaceToBatch(
input_shape,
dilation_rate=dilation_rate,
padding=padding,
build_op=self._build_op,
filter_shape=filter_shape,
spatial_dims=spatial_dims,
data_format=data_format,
)
def from_fields(cls,
fields,
shape=(),
nrows=None,
row_partitions=None,
validate=False):
"""Creates a `StructuredTensor` from a dictionary of fields.
Args:
fields: A dictionary mapping from string to `Tensor`, `RaggedTensor`, or
`StructuredTensor`, providing the values for individual fields in each
structure. If `shape.rank > 0`, then every tensor in `fields` must have
the same shape in the first `shape.rank` dimensions; and that shape must
be compatible with `shape`; and
`result[i1...iN][key] = fields[key][i1...iN]` (where `N==shape.rank`).
shape: A `TensorShape`: static information about the shape of the
`StructuredTensor`. Must have a known `rank`. Defaults to scalar
shape (i.e. `rank=0`).
nrows: scalar integer tensor containing the number of rows in this
`StructuredTensor`. Should only be specified if `shape.rank > 0`.
Default value is inferred from the `fields` values. If `fields` is
empty, then this must be specified.
row_partitions: A list of `RowPartition`s describing the (possibly ragged)
shape of this `StructuredTensor`. Should only be specified if
`shape.rank > 1`. Default value is inferred from the `fields` values.
If `fields` is empty, then this must be specified.
validate: If true, then add runtime validation ops that check that the
field values all have compatible shapes in the outer `shape.rank`
dimensions.
Returns:
A `StructuredTensor`.
Examples:
>>> StructuredTensor.from_fields({'x': 1, 'y': [1, 2, 3]})
<StructuredTensor(fields={
x: tf.Tensor(1, shape=(), dtype=int32),
y: tf.Tensor([1 2 3], shape=(3,), dtype=int32)},
shape=())>
>>> StructuredTensor.from_fields({'foo': [1, 2], 'bar': [3, 4]},
... shape=[2])
<StructuredTensor(fields={
bar: tf.Tensor([3 4], shape=(2,), dtype=int32),
foo: tf.Tensor([1 2], shape=(2,), dtype=int32)},
shape=(2,))>
"""
shape = tensor_shape.as_shape(shape)
rank = shape.rank
if rank is None:
raise ValueError("StructuredTensor's shape must have known rank.")
if not isinstance(fields, dict):
raise TypeError('fields must be a dictionary, got %s' %
type(fields).__name__)
if rank < 2 and row_partitions:
raise ValueError(
'row_partitions must be None or [] if shape.rank<2')
if rank == 0 and nrows is not None:
raise ValueError('nrows must be None if shape.rank==0')
if row_partitions is not None:
row_partitions = tuple(row_partitions)
if len(row_partitions) != max(0, rank - 1):
raise ValueError('len(row_partitions) must be shape.rank-1')
elif rank < 2:
row_partitions = ()
fields = dict(fields) # Make a private copy.
with ops.name_scope(None, 'StructuredTensor', fields.values()):
# Validate keys and convert field values to tensors.
for key, value in fields.items():
if not isinstance(key, str):
raise TypeError('Unexpected type for key in `fields`: %r' %
key)
if not _FIELD_NAME_RE.match(key):
raise ValueError(
'Field name %r is not currently allowed.' % key)
fields[key] = _convert_to_structured_field_value(value)
# Determine dtype for row_partitions and nrows.
shape_dtype = _find_shape_dtype(fields, nrows, row_partitions)
if nrows is not None:
nrows = ops.convert_to_tensor(nrows, shape_dtype)
# Get the static TensorShape for this StructuredTensor.
if rank > 0:
for key, value in fields.items():
if not shape.is_compatible_with(value.shape[:rank]):
raise ValueError(
'Field {} has shape {}, which is incompatible '
'with the shape that was specified or inferred '
'from other fields: {}'.format(
key, value.shape[:rank], shape))
shape = shape.merge_with(value.shape[:rank])
if rank == 1:
# Find a consistent value for `nrows`.
static_nrows = tensor_shape.dimension_at_index(shape, 0)
for value in fields.values():
nrows, static_nrows = _merge_nrows(nrows, static_nrows,
value, shape_dtype,
validate)
if nrows is None:
if static_nrows.value is None:
raise ValueError('nrows must be specified if rank==1 '
'and `fields` is empty.')
else:
nrows = constant_op.constant(static_nrows.value,
shape_dtype)
if rank > 1:
# Find a consistent list of RowPartitions.
for value in fields.values():
row_partitions = _merge_row_partitions(
row_partitions, value, rank, shape_dtype, validate)
if row_partitions is None:
if not shape.is_fully_defined():
raise ValueError(
'row_partitions must be specified if rank>1 '
'and `fields` is empty.')
else:
row_partitions = _row_partitions_for_uniform_shape(
np.array(shape.as_list(),
dtype=shape_dtype.as_numpy_dtype),
shape.rank)
assert len(row_partitions) == rank - 1
nrows = row_partitions[0].nrows()
# Update all field values to use the shared RowPartition objects.
fields = dict([(k, _replace_row_partitions(v, row_partitions))
for (k, v) in fields.items()])
return cls(fields,
shape,
nrows,
row_partitions,
internal=_structured_tensor_factory_key)
