wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/ops/linalg/BUILD
+460
View File
@@ -0,0 +1,460 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
py_library(
name = "linalg_impl",
srcs = ["linalg_impl.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:linalg_ops_gen",
"//tensorflow/python/ops:map_fn",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_util",
srcs = ["linear_operator_util.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_low_rank_update",
srcs = ["linear_operator_low_rank_update.py"],
strict_deps = True,
deps = [
":linear_operator",
":linear_operator_diag",
":linear_operator_identity",
":linear_operator_util",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator_circulant",
srcs = ["linear_operator_circulant.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_util",
":property_hint_util",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops/distributions:util",
"//tensorflow/python/ops/signal:fft_ops",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_identity",
srcs = ["linear_operator_identity.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_diag",
":linear_operator_util",
":property_hint_util",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_kronecker",
srcs = ["linear_operator_kronecker.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
"//tensorflow/python/framework:common_shapes",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "property_hint_util",
srcs = ["property_hint_util.py"],
strict_deps = True,
)
py_library(
name = "linear_operator_block_lower_triangular",
srcs = ["linear_operator_block_lower_triangular.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_addition",
":linear_operator_full_matrix",
":linear_operator_identity",
":linear_operator_util",
"//tensorflow/python/framework:common_shapes",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator",
srcs = [
"linear_operator.py",
"linear_operator_adjoint.py",
"linear_operator_composition.py",
"linear_operator_inversion.py",
"linear_operator_lower_triangular.py",
],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator_util",
":property_hint_util",
":slicing",
"//tensorflow/python/framework:common_shapes",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:composite_tensor_gradient",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/framework:type_spec",
"//tensorflow/python/framework:type_spec_registry",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/trackable:data_structures",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:variable_utils",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_tridiag",
srcs = ["linear_operator_tridiag.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_util",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator_toeplitz",
srcs = ["linear_operator_toeplitz.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_circulant",
":linear_operator_util",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops/signal:fft_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator_block_diag",
srcs = ["linear_operator_block_diag.py"],
strict_deps = True,
deps = [
":linear_operator",
":linear_operator_util",
":property_hint_util",
"//tensorflow/python/framework:common_shapes",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator_test_util",
srcs = ["linear_operator_test_util.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator_util",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:sort_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_v2",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:nested_structure_coder",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_permutation",
srcs = ["linear_operator_permutation.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_util",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sort_ops",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "linalg",
srcs = ["linalg.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_block_diag",
":linear_operator_block_lower_triangular",
":linear_operator_circulant",
":linear_operator_diag",
":linear_operator_full_matrix",
":linear_operator_householder",
":linear_operator_identity",
":linear_operator_kronecker",
":linear_operator_low_rank_update",
":linear_operator_permutation",
":linear_operator_toeplitz",
":linear_operator_tridiag",
":linear_operator_zeros",
],
)
py_library(
name = "linear_operator_full_matrix",
srcs = ["linear_operator_full_matrix.py"],
strict_deps = True,
deps = [
":linear_operator",
":linear_operator_util",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator_householder",
srcs = ["linear_operator_householder.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_util",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "slicing",
srcs = ["slicing.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_zeros",
srcs = ["linear_operator_zeros.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_util",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "linear_operator_diag",
srcs = ["linear_operator_diag.py"],
strict_deps = True,
deps = [
":linalg_impl",
":linear_operator",
":linear_operator_util",
":property_hint_util",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_conversion",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "linear_operator_addition",
srcs = ["linear_operator_addition.py"],
strict_deps = True,
deps = [
":linear_operator",
":linear_operator_diag",
":linear_operator_full_matrix",
":linear_operator_identity",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:check_ops",
],
)
tensorflow/python/ops/linalg/__init__.py
View File
tensorflow/python/ops/linalg/linalg.py
+49
View File
@@ -0,0 +1,49 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Public API for tf.linalg namespace."""
# go/tf-wildcard-import
# pylint: disable=wildcard-import,unused-import
from tensorflow.python.ops.linalg.linalg_impl import *
from tensorflow.python.ops.linalg.linear_operator import *
from tensorflow.python.ops.linalg.linear_operator_adjoint import *
from tensorflow.python.ops.linalg.linear_operator_block_diag import *
from tensorflow.python.ops.linalg.linear_operator_block_lower_triangular import *
from tensorflow.python.ops.linalg.linear_operator_circulant import *
from tensorflow.python.ops.linalg.linear_operator_composition import *
from tensorflow.python.ops.linalg.linear_operator_diag import *
from tensorflow.python.ops.linalg.linear_operator_full_matrix import *
from tensorflow.python.ops.linalg.linear_operator_householder import *
from tensorflow.python.ops.linalg.linear_operator_identity import *
from tensorflow.python.ops.linalg.linear_operator_inversion import *
from tensorflow.python.ops.linalg.linear_operator_kronecker import *
from tensorflow.python.ops.linalg.linear_operator_low_rank_update import *
from tensorflow.python.ops.linalg.linear_operator_lower_triangular import *
from tensorflow.python.ops.linalg.linear_operator_permutation import *
from tensorflow.python.ops.linalg.linear_operator_toeplitz import *
from tensorflow.python.ops.linalg.linear_operator_tridiag import *
from tensorflow.python.ops.linalg.linear_operator_zeros import *
# pylint: enable=wildcard-import
# Seal API.
# pylint: disable=undefined-variable
del ops
del array_ops
del gen_linalg_ops
del linalg_ops
del math_ops
del special_math_ops
del tf_export
# pylint: enable=undefined-variable
tensorflow/python/ops/linalg/linalg_impl.py
+1588
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/linalg/linear_operator.py
+1693
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/linalg/linear_operator_addition.py
+437
View File
@@ -0,0 +1,437 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Add one or more `LinearOperators` efficiently."""
import abc
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_diag
from tensorflow.python.ops.linalg import linear_operator_full_matrix
from tensorflow.python.ops.linalg import linear_operator_identity
from tensorflow.python.ops.linalg import linear_operator_lower_triangular
__all__ = []
def add_operators(operators,
operator_name=None,
addition_tiers=None,
name=None):
"""Efficiently add one or more linear operators.
Given operators `[A1, A2,...]`, this `Op` returns a possibly shorter list of
operators `[B1, B2,...]` such that
```sum_k Ak.matmul(x) = sum_k Bk.matmul(x).```
The operators `Bk` result by adding some of the `Ak`, as allowed by
`addition_tiers`.
Example of efficient adding of diagonal operators.
```python
A1 = LinearOperatorDiag(diag=[1., 1.], name="A1")
A2 = LinearOperatorDiag(diag=[2., 2.], name="A2")
# Use two tiers, the first contains an Adder that returns Diag. Since both
# A1 and A2 are Diag, they can use this Adder. The second tier will not be
# used.
addition_tiers = [
[_AddAndReturnDiag()],
[_AddAndReturnMatrix()]]
B_list = add_operators([A1, A2], addition_tiers=addition_tiers)
len(B_list)
==> 1
B_list[0].__class__.__name__
==> 'LinearOperatorDiag'
B_list[0].to_dense()
==> [[3., 0.],
[0., 3.]]
B_list[0].name
==> 'Add/A1__A2/'
```
Args:
operators: Iterable of `LinearOperator` objects with same `dtype`, domain
and range dimensions, and broadcastable batch shapes.
operator_name: String name for returned `LinearOperator`. Defaults to
concatenation of "Add/A__B/" that indicates the order of addition steps.
addition_tiers: List tiers, like `[tier_0, tier_1, ...]`, where `tier_i`
is a list of `Adder` objects. This function attempts to do all additions
in tier `i` before trying tier `i + 1`.
name: A name for this `Op`. Defaults to `add_operators`.
Returns:
Subclass of `LinearOperator`. Class and order of addition may change as new
(and better) addition strategies emerge.
Raises:
ValueError: If `operators` argument is empty.
ValueError: If shapes are incompatible.
"""
# Default setting
if addition_tiers is None:
addition_tiers = _DEFAULT_ADDITION_TIERS
# Argument checking.
check_ops.assert_proper_iterable(operators)
operators = list(reversed(operators))
if len(operators) < 1:
raise ValueError(
f"Argument `operators` must contain at least one operator. "
f"Received: {operators}.")
if not all(
isinstance(op, linear_operator.LinearOperator) for op in operators):
raise TypeError(
f"Argument `operators` must contain only LinearOperator instances. "
f"Received: {operators}.")
_static_check_for_same_dimensions(operators)
_static_check_for_broadcastable_batch_shape(operators)
with ops.name_scope(name or "add_operators"):
# Additions done in one of the tiers. Try tier 0, 1,...
ops_to_try_at_next_tier = list(operators)
for tier in addition_tiers:
ops_to_try_at_this_tier = ops_to_try_at_next_tier
ops_to_try_at_next_tier = []
while ops_to_try_at_this_tier:
op1 = ops_to_try_at_this_tier.pop()
op2, adder = _pop_a_match_at_tier(op1, ops_to_try_at_this_tier, tier)
if op2 is not None:
# Will try to add the result of this again at this same tier.
new_operator = adder.add(op1, op2, operator_name)
ops_to_try_at_this_tier.append(new_operator)
else:
ops_to_try_at_next_tier.append(op1)
return ops_to_try_at_next_tier
def _pop_a_match_at_tier(op1, operator_list, tier):
# Search from the back of list to the front in order to create nice default
# order of operations.
for i in range(1, len(operator_list) + 1):
op2 = operator_list[-i]
for adder in tier:
if adder.can_add(op1, op2):
return operator_list.pop(-i), adder
return None, None
def _infer_hints_allowing_override(op1, op2, hints):
"""Infer hints from op1 and op2. hints argument is an override.
Args:
op1: LinearOperator
op2: LinearOperator
hints: _Hints object holding "is_X" boolean hints to use for returned
operator.
If some hint is None, try to set using op1 and op2. If the
hint is provided, ignore op1 and op2 hints. This allows an override
of previous hints, but does not allow forbidden hints (e.g. you still
cannot say a real diagonal operator is not self-adjoint.
Returns:
_Hints object.
"""
hints = hints or _Hints()
# If A, B are self-adjoint, then so is A + B.
if hints.is_self_adjoint is None:
is_self_adjoint = op1.is_self_adjoint and op2.is_self_adjoint
else:
is_self_adjoint = hints.is_self_adjoint
# If A, B are positive definite, then so is A + B.
if hints.is_positive_definite is None:
is_positive_definite = op1.is_positive_definite and op2.is_positive_definite
else:
is_positive_definite = hints.is_positive_definite
# A positive definite operator is always non-singular.
if is_positive_definite and hints.is_positive_definite is None:
is_non_singular = True
else:
is_non_singular = hints.is_non_singular
return _Hints(
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite)
def _static_check_for_same_dimensions(operators):
"""ValueError if operators determined to have different dimensions."""
if len(operators) < 2:
return
domain_dimensions = [
(op.name, tensor_shape.dimension_value(op.domain_dimension))
for op in operators
if tensor_shape.dimension_value(op.domain_dimension) is not None]
if len(set(value for name, value in domain_dimensions)) > 1:
raise ValueError(f"All `operators` must have the same `domain_dimension`. "
f"Received: {domain_dimensions}.")
range_dimensions = [
(op.name, tensor_shape.dimension_value(op.range_dimension))
for op in operators
if tensor_shape.dimension_value(op.range_dimension) is not None]
if len(set(value for name, value in range_dimensions)) > 1:
raise ValueError(f"All operators must have the same `range_dimension`. "
f"Received: {range_dimensions}.")
def _static_check_for_broadcastable_batch_shape(operators):
"""ValueError if operators determined to have non-broadcastable shapes."""
if len(operators) < 2:
return
# This will fail if they cannot be broadcast together.
batch_shape = operators[0].batch_shape
for op in operators[1:]:
batch_shape = array_ops.broadcast_static_shape(batch_shape, op.batch_shape)
class _Hints:
"""Holds 'is_X' flags that every LinearOperator is initialized with."""
def __init__(self,
is_non_singular=None,
is_positive_definite=None,
is_self_adjoint=None):
self.is_non_singular = is_non_singular
self.is_positive_definite = is_positive_definite
self.is_self_adjoint = is_self_adjoint
################################################################################
# Classes to add two linear operators.
################################################################################
class _Adder(metaclass=abc.ABCMeta):
"""Abstract base class to add two operators.
Each `Adder` acts independently, adding everything it can, paying no attention
as to whether another `Adder` could have done the addition more efficiently.
"""
@property
def name(self):
return self.__class__.__name__
@abc.abstractmethod
def can_add(self, op1, op2):
"""Returns `True` if this `Adder` can add `op1` and `op2`. Else `False`."""
pass
@abc.abstractmethod
def _add(self, op1, op2, operator_name, hints):
# Derived classes can assume op1 and op2 have been validated, e.g. they have
# the same dtype, and their domain/range dimensions match.
pass
def add(self, op1, op2, operator_name, hints=None):
"""Return new `LinearOperator` acting like `op1 + op2`.
Args:
op1: `LinearOperator`
op2: `LinearOperator`, with `shape` and `dtype` such that adding to
`op1` is allowed.
operator_name: `String` name to give to returned `LinearOperator`
hints: `_Hints` object. Returned `LinearOperator` will be created with
these hints.
Returns:
`LinearOperator`
"""
updated_hints = _infer_hints_allowing_override(op1, op2, hints)
if operator_name is None:
operator_name = "Add/" + op1.name + "__" + op2.name + "/"
scope_name = self.name
if scope_name.startswith("_"):
scope_name = scope_name[1:]
with ops.name_scope(scope_name):
return self._add(op1, op2, operator_name, updated_hints)
class _AddAndReturnScaledIdentity(_Adder):
"""Handles additions resulting in an Identity family member.
The Identity (`LinearOperatorScaledIdentity`, `LinearOperatorIdentity`) family
is closed under addition. This `Adder` respects that, and returns an Identity
"""
def can_add(self, op1, op2):
types = {_type(op1), _type(op2)}
return not types.difference(_IDENTITY_FAMILY)
def _add(self, op1, op2, operator_name, hints):
# Will build a LinearOperatorScaledIdentity.
if _type(op1) == _SCALED_IDENTITY:
multiplier_1 = op1.multiplier
else:
multiplier_1 = array_ops.ones(op1.batch_shape_tensor(), dtype=op1.dtype)
if _type(op2) == _SCALED_IDENTITY:
multiplier_2 = op2.multiplier
else:
multiplier_2 = array_ops.ones(op2.batch_shape_tensor(), dtype=op2.dtype)
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=op1.range_dimension_tensor(),
multiplier=multiplier_1 + multiplier_2,
is_non_singular=hints.is_non_singular,
is_self_adjoint=hints.is_self_adjoint,
is_positive_definite=hints.is_positive_definite,
name=operator_name)
class _AddAndReturnDiag(_Adder):
"""Handles additions resulting in a Diag operator."""
def can_add(self, op1, op2):
types = {_type(op1), _type(op2)}
return not types.difference(_DIAG_LIKE)
def _add(self, op1, op2, operator_name, hints):
return linear_operator_diag.LinearOperatorDiag(
diag=op1.diag_part() + op2.diag_part(),
is_non_singular=hints.is_non_singular,
is_self_adjoint=hints.is_self_adjoint,
is_positive_definite=hints.is_positive_definite,
name=operator_name)
class _AddAndReturnTriL(_Adder):
"""Handles additions resulting in a TriL operator."""
def can_add(self, op1, op2):
types = {_type(op1), _type(op2)}
return not types.difference(_DIAG_LIKE.union({_TRIL}))
def _add(self, op1, op2, operator_name, hints):
if _type(op1) in _EFFICIENT_ADD_TO_TENSOR:
op_add_to_tensor, op_other = op1, op2
else:
op_add_to_tensor, op_other = op2, op1
return linear_operator_lower_triangular.LinearOperatorLowerTriangular(
tril=op_add_to_tensor.add_to_tensor(op_other.to_dense()),
is_non_singular=hints.is_non_singular,
is_self_adjoint=hints.is_self_adjoint,
is_positive_definite=hints.is_positive_definite,
name=operator_name)
class _AddAndReturnMatrix(_Adder):
""""Handles additions resulting in a `LinearOperatorFullMatrix`."""
def can_add(self, op1, op2): # pylint: disable=unused-argument
return isinstance(op1, linear_operator.LinearOperator) and isinstance(
op2, linear_operator.LinearOperator)
def _add(self, op1, op2, operator_name, hints):
if _type(op1) in _EFFICIENT_ADD_TO_TENSOR:
op_add_to_tensor, op_other = op1, op2
else:
op_add_to_tensor, op_other = op2, op1
return linear_operator_full_matrix.LinearOperatorFullMatrix(
matrix=op_add_to_tensor.add_to_tensor(op_other.to_dense()),
is_non_singular=hints.is_non_singular,
is_self_adjoint=hints.is_self_adjoint,
is_positive_definite=hints.is_positive_definite,
name=operator_name)
################################################################################
# Constants designating types of LinearOperators
################################################################################
# Type name constants for LinearOperator classes.
_IDENTITY = "identity"
_SCALED_IDENTITY = "scaled_identity"
_DIAG = "diag"
_TRIL = "tril"
_MATRIX = "matrix"
# Groups of operators.
_DIAG_LIKE = {_DIAG, _IDENTITY, _SCALED_IDENTITY}
_IDENTITY_FAMILY = {_IDENTITY, _SCALED_IDENTITY}
# operators with an efficient .add_to_tensor() method.
_EFFICIENT_ADD_TO_TENSOR = _DIAG_LIKE
# Supported LinearOperator classes.
SUPPORTED_OPERATORS = [
linear_operator_diag.LinearOperatorDiag,
linear_operator_lower_triangular.LinearOperatorLowerTriangular,
linear_operator_full_matrix.LinearOperatorFullMatrix,
linear_operator_identity.LinearOperatorIdentity,
linear_operator_identity.LinearOperatorScaledIdentity
]
def _type(operator):
"""Returns the type name constant (e.g. _TRIL) for operator."""
if isinstance(operator, linear_operator_diag.LinearOperatorDiag):
return _DIAG
if isinstance(operator,
linear_operator_lower_triangular.LinearOperatorLowerTriangular):
return _TRIL
if isinstance(operator, linear_operator_full_matrix.LinearOperatorFullMatrix):
return _MATRIX
if isinstance(operator, linear_operator_identity.LinearOperatorIdentity):
return _IDENTITY
if isinstance(operator,
linear_operator_identity.LinearOperatorScaledIdentity):
return _SCALED_IDENTITY
raise TypeError(f"Expected operator to be one of [LinearOperatorDiag, "
f"LinearOperatorLowerTriangular, LinearOperatorFullMatrix, "
f"LinearOperatorIdentity, LinearOperatorScaledIdentity]. "
f"Received: {operator}")
################################################################################
# Addition tiers:
# We attempt to use Adders in tier K before K+1.
#
# Organize tiers to
# (i) reduce O(..) complexity of forming final operator, and
# (ii) produce the "most efficient" final operator.
# Dev notes:
# * Results of addition at tier K will be added at tier K or higher.
# * Tiers may change, and we warn the user that it may change.
################################################################################
# Note that the final tier, _AddAndReturnMatrix, will convert everything to a
# dense matrix. So it is sometimes very inefficient.
_DEFAULT_ADDITION_TIERS = [
[_AddAndReturnScaledIdentity()],
[_AddAndReturnDiag()],
[_AddAndReturnTriL()],
[_AddAndReturnMatrix()],
]
tensorflow/python/ops/linalg/linear_operator_adjoint.py
+238
View File
@@ -0,0 +1,238 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Takes the adjoint of a `LinearOperator`."""
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorAdjoint"]
@tf_export("linalg.LinearOperatorAdjoint")
@linear_operator.make_composite_tensor
class LinearOperatorAdjoint(linear_operator.LinearOperator):
"""`LinearOperator` representing the adjoint of another operator.
This operator represents the adjoint of another operator.
```python
# Create a 2 x 2 linear operator.
operator = LinearOperatorFullMatrix([[1 - i., 3.], [0., 1. + i]])
operator_adjoint = LinearOperatorAdjoint(operator)
operator_adjoint.to_dense()
==> [[1. + i, 0.]
[3., 1 - i]]
operator_adjoint.shape
==> [2, 2]
operator_adjoint.log_abs_determinant()
==> - log(2)
x = ... Shape [2, 4] Tensor
operator_adjoint.matmul(x)
==> Shape [2, 4] Tensor, equal to operator.matmul(x, adjoint=True)
```
#### Performance
The performance of `LinearOperatorAdjoint` depends on the underlying
operators performance.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
operator,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name=None):
r"""Initialize a `LinearOperatorAdjoint`.
`LinearOperatorAdjoint` is initialized with an operator `A`. The `solve`
and `matmul` methods effectively flip the `adjoint` argument. E.g.
```
A = MyLinearOperator(...)
B = LinearOperatorAdjoint(A)
x = [....] # a vector
assert A.matvec(x, adjoint=True) == B.matvec(x, adjoint=False)
```
Args:
operator: `LinearOperator` object.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`. Default is `operator.name +
"_adjoint"`.
Raises:
ValueError: If `operator.is_non_singular` is False.
"""
parameters = dict(
operator=operator,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name,
)
self._operator = operator
# The congruency of is_non_singular and is_self_adjoint was checked in the
# base operator.
combine_hint = (
linear_operator_util.use_operator_or_provided_hint_unless_contradicting)
is_square = combine_hint(
operator, "is_square", is_square,
"An operator is square if and only if its adjoint is square.")
is_non_singular = combine_hint(
operator, "is_non_singular", is_non_singular,
"An operator is non-singular if and only if its adjoint is "
"non-singular.")
is_self_adjoint = combine_hint(
operator, "is_self_adjoint", is_self_adjoint,
"An operator is self-adjoint if and only if its adjoint is "
"self-adjoint.")
is_positive_definite = combine_hint(
operator, "is_positive_definite", is_positive_definite,
"An operator is positive-definite if and only if its adjoint is "
"positive-definite.")
# Initialization.
if name is None:
name = operator.name + "_adjoint"
with ops.name_scope(name):
super(LinearOperatorAdjoint, self).__init__(
dtype=operator.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
@property
def operator(self):
"""The operator before taking the adjoint."""
return self._operator
def _linop_adjoint(self) -> linear_operator.LinearOperator:
return self.operator
def _assert_non_singular(self):
return self.operator.assert_non_singular()
def _assert_positive_definite(self):
return self.operator.assert_positive_definite()
def _assert_self_adjoint(self):
return self.operator.assert_self_adjoint()
def _shape(self):
# Rotate last dimension
shape = self.operator.shape
return shape[:-2].concatenate([shape[-1], shape[-2]])
def _shape_tensor(self):
# Rotate last dimension
shape = self.operator.shape_tensor()
return array_ops.concat([
shape[:-2], [shape[-1], shape[-2]]], axis=-1)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
return self.operator.matmul(
x, adjoint=(not adjoint), adjoint_arg=adjoint_arg)
def _matvec(self, x, adjoint=False):
return self.operator.matvec(x, adjoint=(not adjoint))
def _determinant(self):
if self.is_self_adjoint:
return self.operator.determinant()
return math_ops.conj(self.operator.determinant())
def _log_abs_determinant(self):
return self.operator.log_abs_determinant()
def _trace(self):
if self.is_self_adjoint:
return self.operator.trace()
return math_ops.conj(self.operator.trace())
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
return self.operator.solve(
rhs, adjoint=(not adjoint), adjoint_arg=adjoint_arg)
def _solvevec(self, rhs, adjoint=False):
return self.operator.solvevec(rhs, adjoint=(not adjoint))
def _to_dense(self):
if self.is_self_adjoint:
return self.operator.to_dense()
return linalg.adjoint(self.operator.to_dense())
def _add_to_tensor(self, x):
return self.to_dense() + x
def _eigvals(self):
eigvals = self.operator.eigvals()
if not self.operator.is_self_adjoint:
eigvals = math_ops.conj(eigvals)
return eigvals
def _cond(self):
return self.operator.cond()
@property
def _composite_tensor_fields(self):
return ("operator",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"operator": 0}
tensorflow/python/ops/linalg/linear_operator_block_diag.py
+818
View File
@@ -0,0 +1,818 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a Block Diagonal operator from one or more `LinearOperators`."""
from tensorflow.python.framework import common_shapes
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.ops.linalg import property_hint_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorBlockDiag"]
@tf_export("linalg.LinearOperatorBlockDiag")
@linear_operator.make_composite_tensor
class LinearOperatorBlockDiag(linear_operator.LinearOperator):
"""Combines one or more `LinearOperators` in to a Block Diagonal matrix.
This operator combines one or more linear operators `[op1,...,opJ]`,
building a new `LinearOperator`, whose underlying matrix representation
has each operator `opi` on the main diagonal, and zero's elsewhere.
#### Shape compatibility
If `opj` acts like a [batch] matrix `Aj`, then `op_combined` acts like
the [batch] matrix formed by having each matrix `Aj` on the main
diagonal.
Each `opj` is required to represent a matrix, and hence will have
shape `batch_shape_j + [M_j, N_j]`.
If `opj` has shape `batch_shape_j + [M_j, N_j]`, then the combined operator
has shape `broadcast_batch_shape + [sum M_j, sum N_j]`, where
`broadcast_batch_shape` is the mutual broadcast of `batch_shape_j`,
`j = 1,...,J`, assuming the intermediate batch shapes broadcast.
Arguments to `matmul`, `matvec`, `solve`, and `solvevec` may either be single
`Tensor`s or lists of `Tensor`s that are interpreted as blocks. The `j`th
element of a blockwise list of `Tensor`s must have dimensions that match
`opj` for the given method. If a list of blocks is input, then a list of
blocks is returned as well.
When the `opj` are not guaranteed to be square, this operator's methods might
fail due to the combined operator not being square and/or lack of efficient
methods.
```python
# Create a 4 x 4 linear operator combined of two 2 x 2 operators.
operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]])
operator_2 = LinearOperatorFullMatrix([[1., 0.], [0., 1.]])
operator = LinearOperatorBlockDiag([operator_1, operator_2])
operator.to_dense()
==> [[1., 2., 0., 0.],
[3., 4., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]]
operator.shape
==> [4, 4]
operator.log_abs_determinant()
==> scalar Tensor
x1 = ... # Shape [2, 2] Tensor
x2 = ... # Shape [2, 2] Tensor
x = tf.concat([x1, x2], 0) # Shape [2, 4] Tensor
operator.matmul(x)
==> tf.concat([operator_1.matmul(x1), operator_2.matmul(x2)])
# Create a 5 x 4 linear operator combining three blocks.
operator_1 = LinearOperatorFullMatrix([[1.], [3.]])
operator_2 = LinearOperatorFullMatrix([[1., 6.]])
operator_3 = LinearOperatorFullMatrix([[2.], [7.]])
operator = LinearOperatorBlockDiag([operator_1, operator_2, operator_3])
operator.to_dense()
==> [[1., 0., 0., 0.],
[3., 0., 0., 0.],
[0., 1., 6., 0.],
[0., 0., 0., 2.]]
[0., 0., 0., 7.]]
operator.shape
==> [5, 4]
# Create a [2, 3] batch of 4 x 4 linear operators.
matrix_44 = tf.random.normal(shape=[2, 3, 4, 4])
operator_44 = LinearOperatorFullMatrix(matrix)
# Create a [1, 3] batch of 5 x 5 linear operators.
matrix_55 = tf.random.normal(shape=[1, 3, 5, 5])
operator_55 = LinearOperatorFullMatrix(matrix_55)
# Combine to create a [2, 3] batch of 9 x 9 operators.
operator_99 = LinearOperatorBlockDiag([operator_44, operator_55])
# Create a shape [2, 3, 9] vector.
x = tf.random.normal(shape=[2, 3, 9])
operator_99.matmul(x)
==> Shape [2, 3, 9] Tensor
# Create a blockwise list of vectors.
x = [tf.random.normal(shape=[2, 3, 4]), tf.random.normal(shape=[2, 3, 5])]
operator_99.matmul(x)
==> [Shape [2, 3, 4] Tensor, Shape [2, 3, 5] Tensor]
```
#### Performance
The performance of `LinearOperatorBlockDiag` on any operation is equal to
the sum of the individual operators' operations.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
operators,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=True,
name=None):
r"""Initialize a `LinearOperatorBlockDiag`.
`LinearOperatorBlockDiag` is initialized with a list of operators
`[op_1,...,op_J]`.
Args:
operators: Iterable of `LinearOperator` objects, each with
the same `dtype` and composable shape.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
This is true by default, and will raise a `ValueError` otherwise.
name: A name for this `LinearOperator`. Default is the individual
operators names joined with `_o_`.
Raises:
TypeError: If all operators do not have the same `dtype`.
ValueError: If `operators` is empty or are non-square.
"""
parameters = dict(
operators=operators,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
# Validate operators.
check_ops.assert_proper_iterable(operators)
operators = list(operators)
if not operators:
raise ValueError(
"Expected a non-empty list of operators. Found: %s" % operators)
self._operators = operators
# Define diagonal operators, for functions that are shared across blockwise
# `LinearOperator` types.
self._diagonal_operators = operators
# Validate dtype.
dtype = operators[0].dtype
for operator in operators:
if operator.dtype != dtype:
name_type = (str((o.name, o.dtype)) for o in operators)
raise TypeError(
"Expected all operators to have the same dtype. Found %s"
% " ".join(name_type))
# Auto-set and check hints.
if all(operator.is_non_singular for operator in operators):
if is_non_singular is False:
raise ValueError(
"The direct sum of non-singular operators is always non-singular.")
is_non_singular = True
if all(operator.is_self_adjoint for operator in operators):
if is_self_adjoint is False:
raise ValueError(
"The direct sum of self-adjoint operators is always self-adjoint.")
is_self_adjoint = True
if all(operator.is_positive_definite for operator in operators):
if is_positive_definite is False:
raise ValueError(
"The direct sum of positive definite operators is always "
"positive definite.")
is_positive_definite = True
if name is None:
# Using ds to mean direct sum.
name = "_ds_".join(operator.name for operator in operators)
with ops.name_scope(name):
super(LinearOperatorBlockDiag, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
@property
def operators(self):
return self._operators
def _block_range_dimensions(self):
return [op.range_dimension for op in self._diagonal_operators]
def _block_domain_dimensions(self):
return [op.domain_dimension for op in self._diagonal_operators]
def _block_range_dimension_tensors(self):
return [op.range_dimension_tensor() for op in self._diagonal_operators]
def _block_domain_dimension_tensors(self):
return [op.domain_dimension_tensor() for op in self._diagonal_operators]
def _shape(self):
# Get final matrix shape.
domain_dimension = sum(self._block_domain_dimensions())
range_dimension = sum(self._block_range_dimensions())
matrix_shape = tensor_shape.TensorShape([range_dimension, domain_dimension])
# Get broadcast batch shape.
# broadcast_shape checks for compatibility.
batch_shape = self.operators[0].batch_shape
for operator in self.operators[1:]:
batch_shape = common_shapes.broadcast_shape(
batch_shape, operator.batch_shape)
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
# Avoid messy broadcasting if possible.
if self.shape.is_fully_defined():
return tensor_conversion.convert_to_tensor_v2_with_dispatch(
self.shape.as_list(), dtype=dtypes.int32, name="shape"
)
domain_dimension = sum(self._block_domain_dimension_tensors())
range_dimension = sum(self._block_range_dimension_tensors())
matrix_shape = array_ops_stack.stack([range_dimension, domain_dimension])
# Dummy Tensor of zeros. Will never be materialized.
zeros = array_ops.zeros(shape=self.operators[0].batch_shape_tensor())
for operator in self.operators[1:]:
zeros += array_ops.zeros(shape=operator.batch_shape_tensor())
batch_shape = array_ops.shape(zeros)
return array_ops.concat((batch_shape, matrix_shape), 0)
def _linop_adjoint(self) -> "LinearOperatorBlockDiag":
# We take the adjoint of each block on the diagonal.
return LinearOperatorBlockDiag(
operators=[operator.adjoint() for operator in self.operators],
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_cholesky(self) -> "LinearOperatorBlockDiag":
# We take the cholesky of each block on the diagonal.
return LinearOperatorBlockDiag(
operators=[operator.cholesky() for operator in self.operators],
is_non_singular=True,
is_self_adjoint=None, # Let the operators passed in decide.
is_square=True)
def _linop_inverse(self) -> "LinearOperatorBlockDiag":
# We take the inverse of each block on the diagonal.
return LinearOperatorBlockDiag(
operators=[
operator.inverse() for operator in self.operators],
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_matmul(
self,
left_operator: "LinearOperatorBlockDiag",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
if isinstance(right_operator, LinearOperatorBlockDiag):
return LinearOperatorBlockDiag(
operators=[
o1.matmul(o2) for o1, o2 in zip(
left_operator.operators, right_operator.operators)],
is_non_singular=property_hint_util.combined_non_singular_hint(
left_operator, right_operator),
# In general, a product of self-adjoint positive-definite
# block diagonal matrices is not self-adjoint.
is_self_adjoint=None,
# In general, a product of positive-definite block diagonal
# matrices is not positive-definite.
is_positive_definite=None,
is_square=True)
return super()._linop_matmul(left_operator, right_operator)
def _linop_solve(
self,
left_operator: "LinearOperatorBlockDiag",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
if isinstance(right_operator, LinearOperatorBlockDiag):
return LinearOperatorBlockDiag(
operators=[
o1.solve(o2) for o1, o2 in zip(
left_operator.operators, right_operator.operators)],
is_non_singular=property_hint_util.combined_non_singular_hint(
left_operator, right_operator),
# In general, a solve of self-adjoint positive-definite block diagonal
# matrices is not self-=adjoint.
is_self_adjoint=None,
# In general, a solve of positive-definite block diagonal matrices is
# not positive-definite.
is_positive_definite=None,
is_square=True)
return super()._linop_solve(left_operator, right_operator)
# TODO(b/188080761): Add a more efficient implementation of `cond` that
# constructs the condition number from the blockwise singular values.
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`, `Tensor` with compatible shape and same `dtype` as
`self`, or a blockwise iterable of `LinearOperator`s or `Tensor`s. See
class docstring for definition of shape 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`, or if `x` is blockwise, a list of `Tensor`s with shapes that
concatenate to `[..., M, R]`.
"""
def _check_operators_agree(r, l, message):
if (r.range_dimension is not None and
l.domain_dimension is not None and
r.range_dimension != l.domain_dimension):
raise ValueError(message)
if isinstance(x, linear_operator.LinearOperator):
left_operator = self.adjoint() if adjoint else self
right_operator = x.adjoint() if adjoint_arg else x
_check_operators_agree(
right_operator, left_operator,
"Operators are incompatible. Expected `x` to have dimension"
" {} but got {}.".format(
left_operator.domain_dimension, right_operator.range_dimension))
# We can efficiently multiply BlockDiag LinearOperators if the number of
# blocks agree.
if isinstance(x, LinearOperatorBlockDiag):
if len(left_operator.operators) != len(right_operator.operators):
raise ValueError(
"Can not efficiently multiply two `LinearOperatorBlockDiag`s "
"together when number of blocks differ.")
for o1, o2 in zip(left_operator.operators, right_operator.operators):
_check_operators_agree(
o2, o1,
"Blocks are incompatible. Expected `x` to have dimension"
" {} but got {}.".format(
o1.domain_dimension, o2.range_dimension))
with self._name_scope(name): # pylint: disable=not-callable
return self._linop_matmul(left_operator, right_operator)
with self._name_scope(name): # pylint: disable=not-callable
arg_dim = -1 if adjoint_arg else -2
block_dimensions = (self._block_range_dimensions() if adjoint
else self._block_domain_dimensions())
if linear_operator_util.arg_is_blockwise(block_dimensions, x, arg_dim):
for i, block in enumerate(x):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[arg_dim])
x[i] = block
else:
x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name="x")
self._check_input_dtype(x)
op_dimension = (self.range_dimension if adjoint
else self.domain_dimension)
op_dimension.assert_is_compatible_with(x.shape[arg_dim])
return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
arg_dim = -1 if adjoint_arg else -2
block_dimensions = (self._block_range_dimensions() if adjoint
else self._block_domain_dimensions())
block_dimensions_fn = (
self._block_range_dimension_tensors if adjoint
else self._block_domain_dimension_tensors)
blockwise_arg = linear_operator_util.arg_is_blockwise(
block_dimensions, x, arg_dim)
if blockwise_arg:
split_x = x
else:
split_dim = -1 if adjoint_arg else -2
# Split input by rows normally, and otherwise columns.
split_x = linear_operator_util.split_arg_into_blocks(
block_dimensions, block_dimensions_fn, x, axis=split_dim)
result_list = []
for index, operator in enumerate(self.operators):
result_list += [operator.matmul(
split_x[index], adjoint=adjoint, adjoint_arg=adjoint_arg)]
if blockwise_arg:
return result_list
result_list = linear_operator_util.broadcast_matrix_batch_dims(
result_list)
return array_ops.concat(result_list, axis=-2)
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`, or an
iterable of `Tensor`s (for blockwise operators). `Tensor`s are treated
a [batch] vectors, 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): # pylint: disable=not-callable
block_dimensions = (self._block_range_dimensions() if adjoint
else self._block_domain_dimensions())
if linear_operator_util.arg_is_blockwise(block_dimensions, x, -1):
for i, block in enumerate(x):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[-1])
x[i] = block
x_mat = [block[..., array_ops.newaxis] for block in x]
y_mat = self.matmul(x_mat, adjoint=adjoint)
return [array_ops.squeeze(y, axis=-1) for y in y_mat]
x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name="x")
self._check_input_dtype(x)
op_dimension = (self.range_dimension if adjoint
else self.domain_dimension)
op_dimension.assert_is_compatible_with(x.shape[-1])
x_mat = x[..., array_ops.newaxis]
y_mat = self.matmul(x_mat, adjoint=adjoint)
return array_ops.squeeze(y_mat, axis=-1)
def _determinant(self):
result = self.operators[0].determinant()
for operator in self.operators[1:]:
result *= operator.determinant()
return result
def _log_abs_determinant(self):
result = self.operators[0].log_abs_determinant()
for operator in self.operators[1:]:
result += operator.log_abs_determinant()
return result
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,
or a list of `Tensor`s (for blockwise operators). `Tensor`s are treated
like a [batch] matrices 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.")
def _check_operators_agree(r, l, message):
if (r.range_dimension is not None and
l.domain_dimension is not None and
r.range_dimension != l.domain_dimension):
raise ValueError(message)
if isinstance(rhs, linear_operator.LinearOperator):
left_operator = self.adjoint() if adjoint else self
right_operator = rhs.adjoint() if adjoint_arg else rhs
_check_operators_agree(
right_operator, left_operator,
"Operators are incompatible. Expected `x` to have dimension"
" {} but got {}.".format(
left_operator.domain_dimension, right_operator.range_dimension))
# We can efficiently solve BlockDiag LinearOperators if the number of
# blocks agree.
if isinstance(right_operator, LinearOperatorBlockDiag):
if len(left_operator.operators) != len(right_operator.operators):
raise ValueError(
"Can not efficiently solve `LinearOperatorBlockDiag` when "
"number of blocks differ.")
for o1, o2 in zip(left_operator.operators, right_operator.operators):
_check_operators_agree(
o2, o1,
"Blocks are incompatible. Expected `x` to have dimension"
" {} but got {}.".format(
o1.domain_dimension, o2.range_dimension))
with self._name_scope(name): # pylint: disable=not-callable
return self._linop_solve(left_operator, right_operator)
with self._name_scope(name): # pylint: disable=not-callable
block_dimensions = (self._block_domain_dimensions() if adjoint
else self._block_range_dimensions())
arg_dim = -1 if adjoint_arg else -2
blockwise_arg = linear_operator_util.arg_is_blockwise(
block_dimensions, rhs, arg_dim)
if blockwise_arg:
split_rhs = rhs
for i, block in enumerate(split_rhs):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[arg_dim])
split_rhs[i] = block
else:
rhs = tensor_conversion.convert_to_tensor_v2_with_dispatch(
rhs, name="rhs"
)
self._check_input_dtype(rhs)
op_dimension = (self.domain_dimension if adjoint
else self.range_dimension)
op_dimension.assert_is_compatible_with(rhs.shape[arg_dim])
split_dim = -1 if adjoint_arg else -2
# Split input by rows normally, and otherwise columns.
split_rhs = linear_operator_util.split_arg_into_blocks(
self._block_domain_dimensions(),
self._block_domain_dimension_tensors,
rhs, axis=split_dim)
solution_list = []
for index, operator in enumerate(self.operators):
solution_list += [operator.solve(
split_rhs[index], adjoint=adjoint, adjoint_arg=adjoint_arg)]
if blockwise_arg:
return solution_list
solution_list = linear_operator_util.broadcast_matrix_batch_dims(
solution_list)
return array_ops.concat(solution_list, axis=-2)
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, or list of `Tensor`s
(for blockwise operators). `Tensor`s are treated as [batch] vectors,
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): # pylint: disable=not-callable
block_dimensions = (self._block_domain_dimensions() if adjoint
else self._block_range_dimensions())
if linear_operator_util.arg_is_blockwise(block_dimensions, rhs, -1):
for i, block in enumerate(rhs):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[-1])
rhs[i] = block
rhs_mat = [array_ops.expand_dims(block, axis=-1) for block in rhs]
solution_mat = self.solve(rhs_mat, adjoint=adjoint)
return [array_ops.squeeze(x, axis=-1) for x in solution_mat]
rhs = tensor_conversion.convert_to_tensor_v2_with_dispatch(
rhs, name="rhs"
)
self._check_input_dtype(rhs)
op_dimension = (self.domain_dimension if adjoint
else self.range_dimension)
op_dimension.assert_is_compatible_with(rhs.shape[-1])
rhs_mat = array_ops.expand_dims(rhs, axis=-1)
solution_mat = self.solve(rhs_mat, adjoint=adjoint)
return array_ops.squeeze(solution_mat, axis=-1)
def _diag_part(self):
if not all(operator.is_square for operator in self.operators):
raise NotImplementedError(
"`diag_part` not implemented for an operator whose blocks are not "
"square.")
diag_list = []
for operator in self.operators:
# Extend the axis for broadcasting.
diag_list += [operator.diag_part()[..., array_ops.newaxis]]
diag_list = linear_operator_util.broadcast_matrix_batch_dims(diag_list)
diagonal = array_ops.concat(diag_list, axis=-2)
return array_ops.squeeze(diagonal, axis=-1)
def _trace(self):
if not all(operator.is_square for operator in self.operators):
raise NotImplementedError(
"`trace` not implemented for an operator whose blocks are not "
"square.")
result = self.operators[0].trace()
for operator in self.operators[1:]:
result += operator.trace()
return result
def _to_dense(self):
num_cols = 0
rows = []
broadcasted_blocks = [operator.to_dense() for operator in self.operators]
broadcasted_blocks = linear_operator_util.broadcast_matrix_batch_dims(
broadcasted_blocks)
for block in broadcasted_blocks:
batch_row_shape = array_ops.shape(block)[:-1]
zeros_to_pad_before_shape = array_ops.concat(
[batch_row_shape, [num_cols]], axis=-1)
zeros_to_pad_before = array_ops.zeros(
shape=zeros_to_pad_before_shape, dtype=block.dtype)
num_cols += array_ops.shape(block)[-1]
zeros_to_pad_after_shape = array_ops.concat(
[batch_row_shape,
[self.domain_dimension_tensor() - num_cols]], axis=-1)
zeros_to_pad_after = array_ops.zeros(
shape=zeros_to_pad_after_shape, dtype=block.dtype)
rows.append(array_ops.concat(
[zeros_to_pad_before, block, zeros_to_pad_after], axis=-1))
mat = array_ops.concat(rows, axis=-2)
mat.set_shape(self.shape)
return mat
def _assert_non_singular(self):
return control_flow_ops.group([
operator.assert_non_singular() for operator in self.operators])
def _assert_self_adjoint(self):
return control_flow_ops.group([
operator.assert_self_adjoint() for operator in self.operators])
def _assert_positive_definite(self):
return control_flow_ops.group([
operator.assert_positive_definite() for operator in self.operators])
def _eigvals(self):
if not all(operator.is_square for operator in self.operators):
raise NotImplementedError(
"`eigvals` not implemented for an operator whose blocks are not "
"square.")
eig_list = []
for operator in self.operators:
# Extend the axis for broadcasting.
eig_list += [operator.eigvals()[..., array_ops.newaxis]]
eig_list = linear_operator_util.broadcast_matrix_batch_dims(eig_list)
eigs = array_ops.concat(eig_list, axis=-2)
return array_ops.squeeze(eigs, axis=-1)
@property
def _composite_tensor_fields(self):
return ("operators",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"operators": [0] * len(self.operators)}
tensorflow/python/ops/linalg/linear_operator_block_lower_triangular.py
+986
View File
@@ -0,0 +1,986 @@
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a blockwise lower-triangular operator from `LinearOperators`."""
from tensorflow.python.framework import common_shapes
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_addition
from tensorflow.python.ops.linalg import linear_operator_full_matrix
from tensorflow.python.ops.linalg import linear_operator_identity
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorBlockLowerTriangular"]
@tf_export("linalg.LinearOperatorBlockLowerTriangular")
@linear_operator.make_composite_tensor
class LinearOperatorBlockLowerTriangular(linear_operator.LinearOperator):
"""Combines `LinearOperators` into a blockwise lower-triangular matrix.
This operator is initialized with a nested list of linear operators, which
are combined into a new `LinearOperator` whose underlying matrix
representation is square and has each operator on or below the main diagonal,
and zero's elsewhere. Each element of the outer list is a list of
`LinearOperators` corresponding to a row-partition of the blockwise structure.
The number of `LinearOperator`s in row-partion `i` must be equal to `i`.
For example, a blockwise `3 x 3` `LinearOperatorBlockLowerTriangular` is
initialized with the list `[[op_00], [op_10, op_11], [op_20, op_21, op_22]]`,
where the `op_ij`, `i < 3, j <= i`, are `LinearOperator` instances. The
`LinearOperatorBlockLowerTriangular` behaves as the following blockwise
matrix, where `0` represents appropriately-sized [batch] matrices of zeros:
```none
[[op_00, 0, 0],
[op_10, op_11, 0],
[op_20, op_21, op_22]]
```
Each `op_jj` on the diagonal is required to represent a square matrix, and
hence will have shape `batch_shape_j + [M_j, M_j]`. `LinearOperator`s in row
`j` of the blockwise structure must have `range_dimension` equal to that of
`op_jj`, and `LinearOperators` in column `j` must have `domain_dimension`
equal to that of `op_jj`.
If each `op_jj` on the diagonal has shape `batch_shape_j + [M_j, M_j]`, then
the combined operator has shape `broadcast_batch_shape + [sum M_j, sum M_j]`,
where `broadcast_batch_shape` is the mutual broadcast of `batch_shape_j`,
`j = 0, 1, ..., J`, assuming the intermediate batch shapes broadcast.
Even if the combined shape is well defined, the combined operator's
methods may fail due to lack of broadcasting ability in the defining
operators' methods.
For example, to create a 4 x 4 linear operator combined of three 2 x 2
operators:
>>> operator_0 = tf.linalg.LinearOperatorFullMatrix([[1., 2.], [3., 4.]])
>>> operator_1 = tf.linalg.LinearOperatorFullMatrix([[1., 0.], [0., 1.]])
>>> operator_2 = tf.linalg.LinearOperatorLowerTriangular([[5., 6.], [7., 8]])
>>> operator = LinearOperatorBlockLowerTriangular(
... [[operator_0], [operator_1, operator_2]])
>>> operator.to_dense()
<tf.Tensor: shape=(4, 4), dtype=float32, numpy=
array([[1., 2., 0., 0.],
[3., 4., 0., 0.],
[1., 0., 5., 0.],
[0., 1., 7., 8.]], dtype=float32)>
>>> operator.shape
TensorShape([4, 4])
>>> operator.log_abs_determinant()
<tf.Tensor: shape=(), dtype=float32, numpy=4.3820267>
>>> x0 = [[1., 6.], [-3., 4.]]
>>> x1 = [[0., 2.], [4., 0.]]
>>> x = tf.concat([x0, x1], 0) # Shape [2, 4] Tensor
>>> operator.matmul(x)
<tf.Tensor: shape=(4, 2), dtype=float32, numpy=
array([[-5., 14.],
[-9., 34.],
[ 1., 16.],
[29., 18.]], dtype=float32)>
The above `matmul` is equivalent to:
>>> tf.concat([operator_0.matmul(x0),
... operator_1.matmul(x0) + operator_2.matmul(x1)], axis=0)
<tf.Tensor: shape=(4, 2), dtype=float32, numpy=
array([[-5., 14.],
[-9., 34.],
[ 1., 16.],
[29., 18.]], dtype=float32)>
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [M, N], with b >= 0
x.shape = [B1,...,Bb] + [N, R], with R >= 0.
```
For example:
Create a [2, 3] batch of 4 x 4 linear operators:
>>> matrix_44 = tf.random.normal(shape=[2, 3, 4, 4])
>>> operator_44 = tf.linalg.LinearOperatorFullMatrix(matrix_44)
Create a [1, 3] batch of 5 x 4 linear operators:
>>> matrix_54 = tf.random.normal(shape=[1, 3, 5, 4])
>>> operator_54 = tf.linalg.LinearOperatorFullMatrix(matrix_54)
Create a [1, 3] batch of 5 x 5 linear operators:
>>> matrix_55 = tf.random.normal(shape=[1, 3, 5, 5])
>>> operator_55 = tf.linalg.LinearOperatorFullMatrix(matrix_55)
Combine to create a [2, 3] batch of 9 x 9 operators:
>>> operator_99 = LinearOperatorBlockLowerTriangular(
... [[operator_44], [operator_54, operator_55]])
>>> operator_99.shape
TensorShape([2, 3, 9, 9])
Create a shape [2, 1, 9] batch of vectors and apply the operator to it.
>>> x = tf.random.normal(shape=[2, 1, 9])
>>> y = operator_99.matvec(x)
>>> y.shape
TensorShape([2, 3, 9])
Create a blockwise list of vectors and apply the operator to it. A blockwise
list is returned.
>>> x4 = tf.random.normal(shape=[2, 1, 4])
>>> x5 = tf.random.normal(shape=[2, 3, 5])
>>> y_blockwise = operator_99.matvec([x4, x5])
>>> y_blockwise[0].shape
TensorShape([2, 3, 4])
>>> y_blockwise[1].shape
TensorShape([2, 3, 5])
#### Performance
Suppose `operator` is a `LinearOperatorBlockLowerTriangular` consisting of `D`
row-partitions and `D` column-partitions, such that the total number of
operators is `N = D * (D + 1) // 2`.
* `operator.matmul` has complexity equal to the sum of the `matmul`
complexities of the individual operators.
* `operator.solve` has complexity equal to the sum of the `solve` complexities
of the operators on the diagonal and the `matmul` complexities of the
operators off the diagonal.
* `operator.determinant` has complexity equal to the sum of the `determinant`
complexities of the operators on the diagonal.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
operators,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorBlockLowerTriangular"):
r"""Initialize a `LinearOperatorBlockLowerTriangular`.
`LinearOperatorBlockLowerTriangular` is initialized with a list of lists of
operators `[[op_0], [op_1, op_2], [op_3, op_4, op_5],...]`.
Args:
operators: Iterable of iterables of `LinearOperator` objects, each with
the same `dtype`. Each element of `operators` corresponds to a row-
partition, in top-to-bottom order. The operators in each row-partition
are filled in left-to-right. For example,
`operators = [[op_0], [op_1, op_2], [op_3, op_4, op_5]]` creates a
`LinearOperatorBlockLowerTriangular` with full block structure
`[[op_0, 0, 0], [op_1, op_2, 0], [op_3, op_4, op_5]]`. The number of
operators in the `i`th row must be equal to `i`, such that each operator
falls on or below the diagonal of the blockwise structure.
`LinearOperator`s that fall on the diagonal (the last elements of each
row) must be square. The other `LinearOperator`s must have domain
dimension equal to the domain dimension of the `LinearOperator`s in the
same column-partition, and range dimension equal to the range dimension
of the `LinearOperator`s in the same row-partition.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
This will raise a `ValueError` if set to `False`.
name: A name for this `LinearOperator`.
Raises:
TypeError: If all operators do not have the same `dtype`.
ValueError: If `operators` is empty, contains an erroneous number of
elements, or contains operators with incompatible shapes.
"""
parameters = dict(
operators=operators,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
# Validate operators.
check_ops.assert_proper_iterable(operators)
for row in operators:
check_ops.assert_proper_iterable(row)
operators = [list(row) for row in operators]
if not operators:
raise ValueError(f"Argument `operators` must be a list of >=1 operators. "
f"Received: {operators}.")
self._operators = operators
self._diagonal_operators = [row[-1] for row in operators]
dtype = operators[0][0].dtype
self._validate_dtype(dtype)
is_non_singular = self._validate_non_singular(is_non_singular)
self._validate_num_operators()
self._validate_operator_dimensions()
is_square = self._validate_square(is_square)
with ops.name_scope(name):
super(LinearOperatorBlockLowerTriangular, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _validate_num_operators(self):
for i, row in enumerate(self.operators):
if len(row) != i + 1:
raise ValueError(
f"Argument `operators[{i}]` must contain `{i + 1}` blocks. "
f"Received: {len(row)} blocks.")
def _validate_operator_dimensions(self):
"""Check that `operators` have compatible dimensions."""
for i in range(1, len(self.operators)):
for j in range(i):
op = self.operators[i][j]
# `above_op` is the operator directly above `op` in the blockwise
# structure, in row partition `i-1`, column partition `j`. `op` should
# have the same `domain_dimension` as `above_op`.
above_op = self.operators[i - 1][j]
# `right_op` is the operator to the right of `op` in the blockwise
# structure, in row partition `i`, column partition `j+1`. `op` should
# have the same `range_dimension` as `right_op`.
right_op = self.operators[i][j + 1]
if (op.domain_dimension is not None and
above_op.domain_dimension is not None):
if op.domain_dimension != above_op.domain_dimension:
raise ValueError(f"Argument `operators[{i}][{j}].domain_dimension` "
f"({op.domain_dimension}) must be the same as "
f"`operators[{i-1}][{j}].domain_dimension` "
f"({above_op.domain_dimension}).")
if (op.range_dimension is not None and
right_op.range_dimension is not None):
if op.range_dimension != right_op.range_dimension:
raise ValueError(f"Argument `operators[{i}][{j}].range_dimension` "
f"({op.range_dimension}) must be the same as "
f"`operators[{i}][{j + 1}].range_dimension` "
f"({right_op.range_dimension}).")
# pylint: disable=g-bool-id-comparison
def _validate_non_singular(self, is_non_singular):
if all(op.is_non_singular for op in self._diagonal_operators):
if is_non_singular is False:
raise ValueError(
f"A blockwise lower-triangular operator with non-singular "
f"operators on the main diagonal is always non-singular. "
f"Expected argument `is_non_singular` to be True. "
f"Received: {is_non_singular}.")
return True
if any(op.is_non_singular is False for op in self._diagonal_operators):
if is_non_singular is True:
raise ValueError(
f"A blockwise lower-triangular operator with a singular operator "
f"on the main diagonal is always singular. Expected argument "
f"`is_non_singular` to be True. Received: {is_non_singular}.")
return False
def _validate_square(self, is_square):
if is_square is False:
raise ValueError(f"`LinearOperatorBlockLowerTriangular` must be square. "
f"Expected argument `is_square` to be True. "
f"Received: {is_square}.")
for i, op in enumerate(self._diagonal_operators):
if op.is_square is False:
raise ValueError(
f"Matrices on the diagonal (the final elements of each "
f"row-partition in the `operators` list) must be square. Expected "
f"argument `operators[{i}][-1].is_square` to be True. "
f"Received: {op.is_square}.")
return True
# pylint: enable=g-bool-id-comparison
def _validate_dtype(self, dtype):
for i, row in enumerate(self.operators):
for operator in row:
if operator.dtype != dtype:
name_type = (str((o.name, o.dtype)) for o in row)
raise TypeError(
"Expected all operators to have the same dtype. Found {} in row "
"{} and {} in row 0.".format(name_type, i, str(dtype)))
@property
def operators(self):
return self._operators
def _block_range_dimensions(self):
return [op.range_dimension for op in self._diagonal_operators]
def _block_domain_dimensions(self):
return [op.domain_dimension for op in self._diagonal_operators]
def _block_range_dimension_tensors(self):
return [op.range_dimension_tensor() for op in self._diagonal_operators]
def _block_domain_dimension_tensors(self):
return [op.domain_dimension_tensor() for op in self._diagonal_operators]
def _shape(self):
# Get final matrix shape.
domain_dimension = sum(self._block_domain_dimensions())
range_dimension = sum(self._block_range_dimensions())
matrix_shape = tensor_shape.TensorShape([domain_dimension, range_dimension])
# Get broadcast batch shape.
# broadcast_shape checks for compatibility.
batch_shape = self.operators[0][0].batch_shape
for row in self.operators[1:]:
for operator in row:
batch_shape = common_shapes.broadcast_shape(
batch_shape, operator.batch_shape)
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
# Avoid messy broadcasting if possible.
if self.shape.is_fully_defined():
return tensor_conversion.convert_to_tensor_v2_with_dispatch(
self.shape.as_list(), dtype=dtypes.int32, name="shape"
)
domain_dimension = sum(self._block_domain_dimension_tensors())
range_dimension = sum(self._block_range_dimension_tensors())
matrix_shape = array_ops_stack.stack([domain_dimension, range_dimension])
batch_shape = self.operators[0][0].batch_shape_tensor()
for row in self.operators[1:]:
for operator in row:
batch_shape = array_ops.broadcast_dynamic_shape(
batch_shape, operator.batch_shape_tensor())
return array_ops.concat((batch_shape, matrix_shape), 0)
def _linop_inverse(self) -> "LinearOperatorBlockLowerTriangular":
"""Inverse of LinearOperatorBlockLowerTriangular.
We recursively apply the identity:
```none
|A 0|' = | A' 0|
|B C| |-C'BA' C'|
```
where `A` is n-by-n, `B` is m-by-n,
`C` is m-by-m, and `'` denotes inverse.
This identity can be verified through multiplication:
```none
|A 0|| A' 0|
|B C||-C'BA' C'|
= | AA' 0|
|BA'-CC'BA' CC'|
= |I 0|
|0 I|
```
Returns:
A 'LinearOperatorBlockLowerTriangular'.
"""
if len(self.operators) == 1:
return (LinearOperatorBlockLowerTriangular(
[[self.operators[0][0].inverse()]],
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=(self.
is_positive_definite),
is_square=True))
blockwise_dim = len(self.operators)
# Calculate the inverse of the `LinearOperatorBlockLowerTriangular`
# representing all but the last row of `self` with
# a recursive call (the matrix `A'` in the docstring definition).
upper_left_inverse = (
LinearOperatorBlockLowerTriangular(self.operators[:-1]).inverse())
bottom_row = self.operators[-1]
bottom_right_inverse = bottom_row[-1].inverse()
# Find the bottom row of the inverse (equal to `[-C'BA', C']`
# in the docstring definition, where `C` is the bottom-right operator of
# `self` and `B` is the set of operators in the
# bottom row excluding `C`). To find `-C'BA'`, we first iterate over the
# column partitions of `A'`.
inverse_bottom_row = []
for i in range(blockwise_dim - 1):
# Find the `i`-th block of `BA'`.
blocks = []
for j in range(i, blockwise_dim - 1):
result = bottom_row[j].matmul(upper_left_inverse.operators[j][i])
if not any(
isinstance(result, op_type)
for op_type in linear_operator_addition.SUPPORTED_OPERATORS
):
result = linear_operator_full_matrix.LinearOperatorFullMatrix(
result.to_dense())
blocks.append(result)
summed_blocks = linear_operator_addition.add_operators(blocks)
assert len(summed_blocks) == 1
block = summed_blocks[0]
# Find the `i`-th block of `-C'BA'`.
block = bottom_right_inverse.matmul(block)
block = linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=bottom_right_inverse.domain_dimension_tensor(),
multiplier=math_ops.cast(-1, dtype=block.dtype)).matmul(block)
inverse_bottom_row.append(block)
# `C'` is the last block of the inverted linear operator.
inverse_bottom_row.append(bottom_right_inverse)
return (LinearOperatorBlockLowerTriangular(
upper_left_inverse.operators + [inverse_bottom_row],
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=(self.is_positive_definite),
is_square=True))
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`, `Tensor` with compatible shape and same `dtype` as
`self`, or a blockwise iterable of `LinearOperator`s or `Tensor`s. See
class docstring for definition of shape 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`, or if `x` is blockwise, a list of `Tensor`s with shapes that
concatenate to `[..., M, R]`.
"""
if isinstance(x, linear_operator.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): # pylint: disable=not-callable
return self._linop_matmul(left_operator, right_operator)
with self._name_scope(name): # pylint: disable=not-callable
arg_dim = -1 if adjoint_arg else -2
block_dimensions = (self._block_range_dimensions() if adjoint
else self._block_domain_dimensions())
if linear_operator_util.arg_is_blockwise(block_dimensions, x, arg_dim):
for i, block in enumerate(x):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[arg_dim])
x[i] = block
else:
x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name="x")
self._check_input_dtype(x)
op_dimension = (self.range_dimension if adjoint
else self.domain_dimension)
op_dimension.assert_is_compatible_with(x.shape[arg_dim])
return self._matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
arg_dim = -1 if adjoint_arg else -2
block_dimensions = (self._block_range_dimensions() if adjoint
else self._block_domain_dimensions())
blockwise_arg = linear_operator_util.arg_is_blockwise(
block_dimensions, x, arg_dim)
if blockwise_arg:
split_x = x
else:
split_dim = -1 if adjoint_arg else -2
# Split input by columns if adjoint_arg is True, else rows
split_x = linear_operator_util.split_arg_into_blocks(
self._block_domain_dimensions(),
self._block_domain_dimension_tensors,
x, axis=split_dim)
result_list = []
# Iterate over row-partitions (i.e. column-partitions of the adjoint).
if adjoint:
for index in range(len(self.operators)):
# Begin with the operator on the diagonal and apply it to the
# respective `rhs` block.
result = self.operators[index][index].matmul(
split_x[index], adjoint=adjoint, adjoint_arg=adjoint_arg)
# Iterate top to bottom over the operators in the remainder of the
# column-partition (i.e. left to right over the row-partition of the
# adjoint), apply the operator to the respective `rhs` block and
# accumulate the sum. For example, given the
# `LinearOperatorBlockLowerTriangular`:
#
# op = [[A, 0, 0],
# [B, C, 0],
# [D, E, F]]
#
# if `index = 1`, the following loop calculates:
# `y_1 = (C.matmul(x_1, adjoint=adjoint) +
# E.matmul(x_2, adjoint=adjoint)`,
# where `x_1` and `x_2` are splits of `x`.
for j in range(index + 1, len(self.operators)):
result += self.operators[j][index].matmul(
split_x[j], adjoint=adjoint, adjoint_arg=adjoint_arg)
result_list.append(result)
else:
for row in self.operators:
# Begin with the left-most operator in the row-partition and apply it
# to the first `rhs` block.
result = row[0].matmul(
split_x[0], adjoint=adjoint, adjoint_arg=adjoint_arg)
# Iterate left to right over the operators in the remainder of the row
# partition, apply the operator to the respective `rhs` block, and
# accumulate the sum.
for j, operator in enumerate(row[1:]):
result += operator.matmul(
split_x[j + 1], adjoint=adjoint, adjoint_arg=adjoint_arg)
result_list.append(result)
if blockwise_arg:
return result_list
result_list = linear_operator_util.broadcast_matrix_batch_dims(
result_list)
return array_ops.concat(result_list, axis=-2)
def matvec(self, x, adjoint=False, name="matvec"):
"""Transform [batch] vector `x` with left multiplication: `x --> Ax`.
```python
# Make an operator acting like batch matrix 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`, or an
iterable of `Tensor`s. `Tensor`s are treated a [batch] vectors, 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): # pylint: disable=not-callable
block_dimensions = (self._block_range_dimensions() if adjoint
else self._block_domain_dimensions())
if linear_operator_util.arg_is_blockwise(block_dimensions, x, -1):
for i, block in enumerate(x):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[-1])
x[i] = block
x_mat = [block[..., array_ops.newaxis] for block in x]
y_mat = self.matmul(x_mat, adjoint=adjoint)
return [array_ops.squeeze(y, axis=-1) for y in y_mat]
x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name="x")
self._check_input_dtype(x)
op_dimension = (self.range_dimension if adjoint
else self.domain_dimension)
op_dimension.assert_is_compatible_with(x.shape[-1])
x_mat = x[..., array_ops.newaxis]
y_mat = self.matmul(x_mat, adjoint=adjoint)
return array_ops.squeeze(y_mat, axis=-1)
def _determinant(self):
if all(op.is_positive_definite for op in self._diagonal_operators):
return math_ops.exp(self._log_abs_determinant())
result = self._diagonal_operators[0].determinant()
for op in self._diagonal_operators[1:]:
result *= op.determinant()
return result
def _log_abs_determinant(self):
result = self._diagonal_operators[0].log_abs_determinant()
for op in self._diagonal_operators[1:]:
result += op.log_abs_determinant()
return result
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.
Given the blockwise `n + 1`-by-`n + 1` linear operator:
op = [[A_00 0 ... 0 ... 0],
[A_10 A_11 ... 0 ... 0],
...
[A_k0 A_k1 ... A_kk ... 0],
...
[A_n0 A_n1 ... A_nk ... A_nn]]
we find `x = op.solve(y)` by observing that
`y_k = A_k0.matmul(x_0) + A_k1.matmul(x_1) + ... + A_kk.matmul(x_k)`
and therefore
`x_k = A_kk.solve(y_k -
A_k0.matmul(x_0) - ... - A_k(k-1).matmul(x_(k-1)))`
where `x_k` and `y_k` are the `k`th blocks obtained by decomposing `x`
and `y` along their appropriate axes.
We first solve `x_0 = A_00.solve(y_0)`. Proceeding inductively, we solve
for `x_k`, `k = 1..n`, given `x_0..x_(k-1)`.
The adjoint case is solved similarly, beginning with
`x_n = A_nn.solve(y_n, adjoint=True)` and proceeding backwards.
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,
or a list of `Tensor`s. `Tensor`s are treated like a [batch] matrices
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.")
if isinstance(rhs, linear_operator.LinearOperator):
left_operator = self.adjoint() if adjoint else self
right_operator = rhs.adjoint() if adjoint_arg else rhs
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 `rhs` to have dimension"
" {} but got {}.".format(
left_operator.domain_dimension, right_operator.range_dimension))
with self._name_scope(name): # pylint: disable=not-callable
return self._linop_solve(left_operator, right_operator)
with self._name_scope(name): # pylint: disable=not-callable
block_dimensions = (self._block_domain_dimensions() if adjoint
else self._block_range_dimensions())
arg_dim = -1 if adjoint_arg else -2
blockwise_arg = linear_operator_util.arg_is_blockwise(
block_dimensions, rhs, arg_dim)
if blockwise_arg:
for i, block in enumerate(rhs):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[arg_dim])
rhs[i] = block
if adjoint_arg:
split_rhs = [linalg.adjoint(y) for y in rhs]
else:
split_rhs = rhs
else:
rhs = tensor_conversion.convert_to_tensor_v2_with_dispatch(
rhs, name="rhs"
)
self._check_input_dtype(rhs)
op_dimension = (self.domain_dimension if adjoint
else self.range_dimension)
op_dimension.assert_is_compatible_with(rhs.shape[arg_dim])
rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
split_rhs = linear_operator_util.split_arg_into_blocks(
self._block_domain_dimensions(),
self._block_domain_dimension_tensors,
rhs, axis=-2)
solution_list = []
if adjoint:
# For an adjoint blockwise lower-triangular linear operator, the system
# must be solved bottom to top. Iterate backwards over rows of the
# adjoint (i.e. columns of the non-adjoint operator).
for index in reversed(range(len(self.operators))):
y = split_rhs[index]
# Iterate top to bottom over the operators in the off-diagonal portion
# of the column-partition (i.e. row-partition of the adjoint), apply
# the operator to the respective block of the solution found in
# previous iterations, and subtract the result from the `rhs` block.
# For example,let `A`, `B`, and `D` be the linear operators in the top
# row-partition of the adjoint of
# `LinearOperatorBlockLowerTriangular([[A], [B, C], [D, E, F]])`,
# and `x_1` and `x_2` be blocks of the solution found in previous
# iterations of the outer loop. The following loop (when `index == 0`)
# expresses
# `Ax_0 + Bx_1 + Dx_2 = y_0` as `Ax_0 = y_0*`, where
# `y_0* = y_0 - Bx_1 - Dx_2`.
for j in reversed(range(index + 1, len(self.operators))):
y = y - self.operators[j][index].matmul(
solution_list[len(self.operators) - 1 - j],
adjoint=adjoint)
# Continuing the example above, solve `Ax_0 = y_0*` for `x_0`.
solution_list.append(
self._diagonal_operators[index].solve(y, adjoint=adjoint))
solution_list.reverse()
else:
# Iterate top to bottom over the row-partitions.
for row, y in zip(self.operators, split_rhs):
# Iterate left to right over the operators in the off-diagonal portion
# of the row-partition, apply the operator to the block of the
# solution found in previous iterations, and subtract the result from
# the `rhs` block. For example, let `D`, `E`, and `F` be the linear
# operators in the bottom row-partition of
# `LinearOperatorBlockLowerTriangular([[A], [B, C], [D, E, F]])` and
# `x_0` and `x_1` be blocks of the solution found in previous
# iterations of the outer loop. The following loop
# (when `index == 2`), expresses
# `Dx_0 + Ex_1 + Fx_2 = y_2` as `Fx_2 = y_2*`, where
# `y_2* = y_2 - D_x0 - Ex_1`.
for i, operator in enumerate(row[:-1]):
y = y - operator.matmul(solution_list[i], adjoint=adjoint)
# Continuing the example above, solve `Fx_2 = y_2*` for `x_2`.
solution_list.append(row[-1].solve(y, adjoint=adjoint))
if blockwise_arg:
return solution_list
solution_list = linear_operator_util.broadcast_matrix_batch_dims(
solution_list)
return array_ops.concat(solution_list, axis=-2)
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, or list of `Tensor`s
(for blockwise operators). `Tensor`s are treated as [batch] vectors,
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): # pylint: disable=not-callable
block_dimensions = (self._block_domain_dimensions() if adjoint
else self._block_range_dimensions())
if linear_operator_util.arg_is_blockwise(block_dimensions, rhs, -1):
for i, block in enumerate(rhs):
if not isinstance(block, linear_operator.LinearOperator):
block = tensor_conversion.convert_to_tensor_v2_with_dispatch(block)
self._check_input_dtype(block)
block_dimensions[i].assert_is_compatible_with(block.shape[-1])
rhs[i] = block
rhs_mat = [array_ops.expand_dims(block, axis=-1) for block in rhs]
solution_mat = self.solve(rhs_mat, adjoint=adjoint)
return [array_ops.squeeze(x, axis=-1) for x in solution_mat]
rhs = tensor_conversion.convert_to_tensor_v2_with_dispatch(
rhs, name="rhs"
)
self._check_input_dtype(rhs)
op_dimension = (self.domain_dimension if adjoint
else self.range_dimension)
op_dimension.assert_is_compatible_with(rhs.shape[-1])
rhs_mat = array_ops.expand_dims(rhs, axis=-1)
solution_mat = self.solve(rhs_mat, adjoint=adjoint)
return array_ops.squeeze(solution_mat, axis=-1)
def _diag_part(self):
diag_list = []
for op in self._diagonal_operators:
# Extend the axis, since `broadcast_matrix_batch_dims` treats all but the
# final two dimensions as batch dimensions.
diag_list.append(op.diag_part()[..., array_ops.newaxis])
diag_list = linear_operator_util.broadcast_matrix_batch_dims(diag_list)
diagonal = array_ops.concat(diag_list, axis=-2)
return array_ops.squeeze(diagonal, axis=-1)
def _trace(self):
result = self._diagonal_operators[0].trace()
for op in self._diagonal_operators[1:]:
result += op.trace()
return result
def _to_dense(self):
num_cols = 0
dense_rows = []
flat_broadcast_operators = linear_operator_util.broadcast_matrix_batch_dims(
[op.to_dense() for row in self.operators for op in row]) # pylint: disable=g-complex-comprehension
broadcast_operators = [
flat_broadcast_operators[i * (i + 1) // 2:(i + 1) * (i + 2) // 2]
for i in range(len(self.operators))]
for row_blocks in broadcast_operators:
batch_row_shape = array_ops.shape(row_blocks[0])[:-1]
num_cols += array_ops.shape(row_blocks[-1])[-1]
zeros_to_pad_after_shape = array_ops.concat(
[batch_row_shape,
[self.domain_dimension_tensor() - num_cols]], axis=-1)
zeros_to_pad_after = array_ops.zeros(
shape=zeros_to_pad_after_shape, dtype=self.dtype)
row_blocks.append(zeros_to_pad_after)
dense_rows.append(array_ops.concat(row_blocks, axis=-1))
mat = array_ops.concat(dense_rows, axis=-2)
mat.set_shape(self.shape)
return mat
def _assert_non_singular(self):
return control_flow_ops.group([
op.assert_non_singular() for op in self._diagonal_operators])
def _eigvals(self):
eig_list = []
for op in self._diagonal_operators:
# Extend the axis for broadcasting.
eig_list.append(op.eigvals()[..., array_ops.newaxis])
eig_list = linear_operator_util.broadcast_matrix_batch_dims(eig_list)
eigs = array_ops.concat(eig_list, axis=-2)
return array_ops.squeeze(eigs, axis=-1)
@property
def _composite_tensor_fields(self):
return ("operators",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
# None of the operators contribute to the matrix shape.
return {"operators": nest.map_structure(lambda _: 0, self.operators)}
tensorflow/python/ops/linalg/linear_operator_circulant.py
+1557
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/linalg/linear_operator_composition.py
+404
View File
@@ -0,0 +1,404 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Composes one or more `LinearOperators`."""
from tensorflow.python.framework import common_shapes
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_lower_triangular
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorComposition"]
@tf_export("linalg.LinearOperatorComposition")
@linear_operator.make_composite_tensor
class LinearOperatorComposition(linear_operator.LinearOperator):
"""Composes one or more `LinearOperators`.
This operator composes one or more linear operators `[op1,...,opJ]`,
building a new `LinearOperator` with action defined by:
```
op_composed(x) := op1(op2(...(opJ(x)...))
```
If `opj` acts like [batch] matrix `Aj`, then `op_composed` acts like the
[batch] matrix formed with the multiplication `A1 A2...AJ`.
If `opj` has shape `batch_shape_j + [M_j, N_j]`, then we must have
`N_j = M_{j+1}`, in which case the composed operator has shape equal to
`broadcast_batch_shape + [M_1, N_J]`, where `broadcast_batch_shape` is the
mutual broadcast of `batch_shape_j`, `j = 1,...,J`, assuming the intermediate
batch shapes broadcast. Even if the composed shape is well defined, the
composed operator's methods may fail due to lack of broadcasting ability in
the defining operators' methods.
```python
# Create a 2 x 2 linear operator composed of two 2 x 2 operators.
operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]])
operator_2 = LinearOperatorFullMatrix([[1., 0.], [0., 1.]])
operator = LinearOperatorComposition([operator_1, operator_2])
operator.to_dense()
==> [[1., 2.]
[3., 4.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor
# Create a [2, 3] batch of 4 x 5 linear operators.
matrix_45 = tf.random.normal(shape=[2, 3, 4, 5])
operator_45 = LinearOperatorFullMatrix(matrix)
# Create a [2, 3] batch of 5 x 6 linear operators.
matrix_56 = tf.random.normal(shape=[2, 3, 5, 6])
operator_56 = LinearOperatorFullMatrix(matrix_56)
# Compose to create a [2, 3] batch of 4 x 6 operators.
operator_46 = LinearOperatorComposition([operator_45, operator_56])
# Create a shape [2, 3, 6, 2] vector.
x = tf.random.normal(shape=[2, 3, 6, 2])
operator.matmul(x)
==> Shape [2, 3, 4, 2] Tensor
```
#### Performance
The performance of `LinearOperatorComposition` on any operation is equal to
the sum of the individual operators' operations.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
operators,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name=None):
r"""Initialize a `LinearOperatorComposition`.
`LinearOperatorComposition` is initialized with a list of operators
`[op_1,...,op_J]`. For the `matmul` method to be well defined, the
composition `op_i.matmul(op_{i+1}(x))` must be defined. Other methods have
similar constraints.
Args:
operators: Iterable of `LinearOperator` objects, each with
the same `dtype` and composable shape.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`. Default is the individual
operators names joined with `_o_`.
Raises:
TypeError: If all operators do not have the same `dtype`.
ValueError: If `operators` is empty.
"""
parameters = dict(
operators=operators,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name)
# Validate operators.
check_ops.assert_proper_iterable(operators)
operators = list(operators)
if not operators:
raise ValueError(
"Expected a non-empty list of operators. Found: %s" % operators)
self._operators = operators
# Validate dtype.
dtype = operators[0].dtype
for operator in operators:
if operator.dtype != dtype:
name_type = (str((o.name, o.dtype)) for o in operators)
raise TypeError(
"Expected all operators to have the same dtype. Found %s"
% " ".join(name_type))
# Auto-set and check hints.
if all(operator.is_non_singular for operator in operators):
if is_non_singular is False: # pylint:disable=g-bool-id-comparison
raise ValueError(
"The composition of non-singular operators is always non-singular.")
is_non_singular = True
if _composition_must_be_self_adjoint(operators):
if is_self_adjoint is False: # pylint:disable=g-bool-id-comparison
raise ValueError(
"The composition was determined to be self-adjoint but user "
"provided incorrect `False` hint.")
is_self_adjoint = True
if linear_operator_util.is_aat_form(operators):
if is_square is False: # pylint:disable=g-bool-id-comparison
raise ValueError(
"The composition was determined have the form "
"A @ A.H, hence it must be square. The user "
"provided an incorrect `False` hint.")
is_square = True
if linear_operator_util.is_aat_form(operators) and is_non_singular:
if is_positive_definite is False: # pylint:disable=g-bool-id-comparison
raise ValueError(
"The composition was determined to be non-singular and have the "
"form A @ A.H, hence it must be positive-definite. The user "
"provided an incorrect `False` hint.")
is_positive_definite = True
# Initialization.
if name is None:
name = "_o_".join(operator.name for operator in operators)
with ops.name_scope(name):
super(LinearOperatorComposition, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
@property
def operators(self):
return self._operators
def _shape(self):
# Get final matrix shape.
domain_dimension = self.operators[0].domain_dimension
for operator in self.operators[1:]:
domain_dimension.assert_is_compatible_with(operator.range_dimension)
domain_dimension = operator.domain_dimension
matrix_shape = tensor_shape.TensorShape(
[self.operators[0].range_dimension,
self.operators[-1].domain_dimension])
# Get broadcast batch shape.
# broadcast_shape checks for compatibility.
batch_shape = self.operators[0].batch_shape
for operator in self.operators[1:]:
batch_shape = common_shapes.broadcast_shape(
batch_shape, operator.batch_shape)
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
# Avoid messy broadcasting if possible.
if self.shape.is_fully_defined():
return ops.convert_to_tensor(
self.shape.as_list(), dtype=dtypes.int32, name="shape")
# Don't check the matrix dimensions. That would add unnecessary Asserts to
# the graph. Things will fail at runtime naturally if shapes are
# incompatible.
matrix_shape = array_ops_stack.stack([
self.operators[0].range_dimension_tensor(),
self.operators[-1].domain_dimension_tensor()
])
# Dummy Tensor of zeros. Will never be materialized.
zeros = array_ops.zeros(shape=self.operators[0].batch_shape_tensor())
for operator in self.operators[1:]:
zeros += array_ops.zeros(shape=operator.batch_shape_tensor())
batch_shape = array_ops.shape(zeros)
return array_ops.concat((batch_shape, matrix_shape), 0)
def _linop_cholesky(self) -> linear_operator.LinearOperator:
"""Computes Cholesky(LinearOperatorComposition)."""
# L @ L.H will be handled with special code below. Why is L @ L.H the most
# important special case?
# Note that Diag @ Diag.H and Diag @ TriL and TriL @ Diag are already
# compressed to Diag or TriL by diag matmul
# registration. Similarly for Identity and ScaledIdentity.
# So these would not appear in a LinearOperatorComposition unless explicitly
# constructed as such. So the most important thing to check is L @ L.H.
def _is_llt_product(self):
"""Determines if linop = L @ L.H for L = LinearOperatorLowerTriangular."""
if len(self.operators) != 2:
return False
if not linear_operator_util.is_aat_form(self.operators):
return False
return isinstance(
self.operators[0],
linear_operator_lower_triangular.LinearOperatorLowerTriangular)
if not _is_llt_product(self):
return linear_operator_lower_triangular.LinearOperatorLowerTriangular(
linalg_ops.cholesky(self.to_dense()),
is_non_singular=True,
is_self_adjoint=False,
is_square=True)
left_op = self.operators[0]
# left_op.is_positive_definite ==> op already has positive diag,return it.
if left_op.is_positive_definite:
return left_op
# Recall that the base class has already verified
# linop.is_positive_definite, else linop.cholesky() would have raised.
# So in particular, we know the diagonal has nonzero entries.
# In the generic case, we make op have positive diag by dividing each row
# by the sign of the diag. This is equivalent to setting A = L @ D where
# D is diag(sign(1 / L.diag_part())). Then A is lower triangular with
# positive diag and A @ A^H = L @ D @ D^H @ L^H = L @ L^H = linop.
# This also works for complex L,
# since sign(x + iy) = exp(i * angle(x + iy)).
diag_sign = array_ops.expand_dims(
math_ops.sign(left_op.diag_part()), axis=-2)
return linear_operator_lower_triangular.LinearOperatorLowerTriangular(
tril=left_op.tril / diag_sign,
is_non_singular=left_op.is_non_singular,
# L.is_self_adjoint ==> L is diagonal ==> L @ D is diagonal ==> SA
# L.is_self_adjoint is False ==> L not diagonal ==> L @ D not diag ...
is_self_adjoint=left_op.is_self_adjoint,
# L.is_positive_definite ==> L has positive diag ==> L = L @ D
# ==> (L @ D).is_positive_definite.
# L.is_positive_definite is False could result
# in L @ D being PD or not.
# Consider L = [[1, 0], [-2, 1]] and quadratic form with x = [1, 1].
# Note we will already return left_op if left_op.is_positive_definite
# above, but to be explicit write this below.
is_positive_definite=True if left_op.is_positive_definite else None,
is_square=True,
)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
# If self.operators = [A, B], and not adjoint, then
# matmul_order_list = [B, A].
# As a result, we return A.matmul(B.matmul(x))
if adjoint:
matmul_order_list = self.operators
else:
matmul_order_list = list(reversed(self.operators))
result = matmul_order_list[0].matmul(
x, adjoint=adjoint, adjoint_arg=adjoint_arg)
for operator in matmul_order_list[1:]:
result = operator.matmul(result, adjoint=adjoint)
return result
def _determinant(self):
result = self.operators[0].determinant()
for operator in self.operators[1:]:
result *= operator.determinant()
return result
def _log_abs_determinant(self):
result = self.operators[0].log_abs_determinant()
for operator in self.operators[1:]:
result += operator.log_abs_determinant()
return result
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
# TODO(langmore) Implement solve using solve_ls if some intermediate
# operator maps to a high dimensional space.
# In that case, an exact solve may still be possible.
# If self.operators = [A, B], and not adjoint, then
# solve_order_list = [A, B].
# As a result, we return B.solve(A.solve(x))
if adjoint:
solve_order_list = list(reversed(self.operators))
else:
solve_order_list = self.operators
solution = solve_order_list[0].solve(
rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
for operator in solve_order_list[1:]:
solution = operator.solve(solution, adjoint=adjoint)
return solution
def _assert_non_singular(self):
if all(operator.is_square for operator in self.operators):
asserts = [operator.assert_non_singular() for operator in self.operators]
return control_flow_ops.group(asserts)
return super(LinearOperatorComposition, self)._assert_non_singular()
@property
def _composite_tensor_fields(self):
return ("operators",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"operators": [0] * len(self.operators)}
def _composition_must_be_self_adjoint(operators):
"""Runs some checks to see if composition operators must be SA.
Args:
operators: List of LinearOperators.
Returns:
True if the composition must be SA. False if it is not SA OR if we did not
determine whether the composition is SA.
"""
if len(operators) == 1 and operators[0].is_self_adjoint:
return True
# Check for forms like A @ A.H or (A1 @ A2) @ (A2.H @ A1.H) or ...
if linear_operator_util.is_aat_form(operators):
return True
# Done checking...could still be SA.
# We may not catch some cases. E.g. (A @ I) @ A.H is SA, but is not AAT form.
return False
tensorflow/python/ops/linalg/linear_operator_diag.py
+388
View File
@@ -0,0 +1,388 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a diagonal matrix."""
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_lower_triangular
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.ops.linalg import property_hint_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorDiag",]
@tf_export("linalg.LinearOperatorDiag")
@linear_operator.make_composite_tensor
class LinearOperatorDiag(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] square diagonal matrix.
This operator acts like a [batch] diagonal matrix `A` with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x N` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
`LinearOperatorDiag` is initialized with a (batch) vector.
```python
# Create a 2 x 2 diagonal linear operator.
diag = [1., -1.]
operator = LinearOperatorDiag(diag)
operator.to_dense()
==> [[1., 0.]
[0., -1.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor
# Create a [2, 3] batch of 4 x 4 linear operators.
diag = tf.random.normal(shape=[2, 3, 4])
operator = LinearOperatorDiag(diag)
# Create a shape [2, 1, 4, 2] vector. Note that this shape is compatible
# since the batch dimensions, [2, 1], are broadcast to
# operator.batch_shape = [2, 3].
y = tf.random.normal(shape=[2, 1, 4, 2])
x = operator.solve(y)
==> operator.matmul(x) = y
```
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
#### Performance
Suppose `operator` is a `LinearOperatorDiag` of shape `[N, N]`,
and `x.shape = [N, R]`. Then
* `operator.matmul(x)` involves `N * R` multiplications.
* `operator.solve(x)` involves `N` divisions and `N * R` multiplications.
* `operator.determinant()` involves a size `N` `reduce_prod`.
If instead `operator` and `x` have shape `[B1,...,Bb, N, N]` and
`[B1,...,Bb, N, R]`, every operation increases in complexity by `B1*...*Bb`.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
diag,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorDiag"):
r"""Initialize a `LinearOperatorDiag`.
Args:
diag: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`.
The diagonal of the operator. Allowed dtypes: `float16`, `float32`,
`float64`, `complex64`, `complex128`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. If `diag.dtype` is real, this is auto-set to `True`.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`.
Raises:
TypeError: If `diag.dtype` is not an allowed type.
ValueError: If `diag.dtype` is real, and `is_self_adjoint` is not `True`.
"""
parameters = dict(
diag=diag,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
with ops.name_scope(name, values=[diag]):
self._diag = linear_operator_util.convert_nonref_to_tensor(
diag, name="diag")
self._check_diag(self._diag)
# Check and auto-set hints.
if not self._diag.dtype.is_complex:
if is_self_adjoint is False:
raise ValueError("A real diagonal operator is always self adjoint.")
else:
is_self_adjoint = True
if is_square is False:
raise ValueError("Only square diagonal operators currently supported.")
is_square = True
super(LinearOperatorDiag, self).__init__(
dtype=self._diag.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _check_diag(self, diag):
"""Static check of diag."""
if diag.shape.ndims is not None and diag.shape.ndims < 1:
raise ValueError("Argument diag must have at least 1 dimension. "
"Found: %s" % diag)
def _shape(self):
# If d_shape = [5, 3], we return [5, 3, 3].
d_shape = self._diag.shape
return d_shape.concatenate(d_shape[-1:])
def _shape_tensor(self):
d_shape = array_ops.shape(self._diag)
k = d_shape[-1]
return array_ops.concat((d_shape, [k]), 0)
@property
def diag(self):
return self._diag
def _linop_inverse(self) -> "LinearOperatorDiag":
return LinearOperatorDiag(
1. / self.diag,
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_matmul(
self,
left_operator: "LinearOperatorDiag",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
is_non_singular = property_hint_util.combined_non_singular_hint(
left_operator, right_operator)
is_self_adjoint = property_hint_util.combined_commuting_self_adjoint_hint(
left_operator, right_operator)
is_positive_definite = (
property_hint_util.combined_commuting_positive_definite_hint(
left_operator, right_operator))
if isinstance(right_operator, LinearOperatorDiag):
return LinearOperatorDiag(
diag=left_operator.diag * right_operator.diag,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True,
)
# instance of linear_operator_identity.LinearOperatorScaledIdentity
elif hasattr(right_operator, "_ones_diag") and hasattr(
right_operator, "multiplier"
):
return LinearOperatorDiag(
diag=left_operator.diag * right_operator.multiplier,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
elif isinstance(
right_operator,
linear_operator_lower_triangular.LinearOperatorLowerTriangular,
):
return linear_operator_lower_triangular.LinearOperatorLowerTriangular(
tril=left_operator.diag[..., None] * right_operator.to_dense(),
is_non_singular=is_non_singular,
# This is safe to do since the Triangular matrix is only self-adjoint
# when it is a diagonal matrix, and hence commutes.
is_self_adjoint=is_self_adjoint,
is_positive_definite=None,
is_square=True)
else:
return super()._linop_matmul(left_operator, right_operator)
def _linop_solve(
self,
left_operator: "LinearOperatorDiag",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
is_non_singular = property_hint_util.combined_non_singular_hint(
left_operator, right_operator)
is_self_adjoint = property_hint_util.combined_commuting_self_adjoint_hint(
left_operator, right_operator)
is_positive_definite = (
property_hint_util.combined_commuting_positive_definite_hint(
left_operator, right_operator))
if isinstance(right_operator, LinearOperatorDiag):
return LinearOperatorDiag(
diag=right_operator.diag / left_operator.diag,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
# instance of linear_operator_identity.LinearOperatorScaledIdentity
elif (hasattr(right_operator, "_ones_diag")
and hasattr(right_operator, "multiplier")):
return LinearOperatorDiag(
diag=right_operator.multiplier / left_operator.diag,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
elif isinstance(
right_operator,
linear_operator_lower_triangular.LinearOperatorLowerTriangular):
return linear_operator_lower_triangular.LinearOperatorLowerTriangular(
tril=right_operator.to_dense() / left_operator.diag[..., None],
is_non_singular=is_non_singular,
# This is safe to do since the Triangular matrix is only self-adjoint
# when it is a diagonal matrix, and hence commutes.
is_self_adjoint=is_self_adjoint,
is_positive_definite=None,
is_square=True)
else:
return super()._linop_solve(left_operator, right_operator)
def _assert_non_singular(self):
return linear_operator_util.assert_no_entries_with_modulus_zero(
self._diag,
message="Singular operator: Diagonal contained zero values.")
def _assert_positive_definite(self):
if self.dtype.is_complex:
message = (
"Diagonal operator had diagonal entries with non-positive real part, "
"thus was not positive definite.")
else:
message = (
"Real diagonal operator had non-positive diagonal entries, "
"thus was not positive definite.")
return check_ops.assert_positive(
math_ops.real(self._diag),
message=message)
def _assert_self_adjoint(self):
return linear_operator_util.assert_zero_imag_part(
self._diag,
message=(
"This diagonal operator contained non-zero imaginary values. "
" Thus it was not self-adjoint."))
def _linop_adjoint(self) -> "LinearOperatorDiag":
diag = self.diag
if diag.dtype.is_complex:
diag = math_ops.conj(diag)
return LinearOperatorDiag(
diag=diag,
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_cholesky(self) -> "LinearOperatorDiag":
return LinearOperatorDiag(
math_ops.sqrt(self.diag),
is_non_singular=True,
is_self_adjoint=True,
is_positive_definite=True,
is_square=True)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
diag_term = math_ops.conj(self._diag) if adjoint else self._diag
x = linalg.adjoint(x) if adjoint_arg else x
diag_mat = array_ops.expand_dims(diag_term, -1)
return diag_mat * x
def _matvec(self, x, adjoint=False):
diag_term = math_ops.conj(self._diag) if adjoint else self._diag
return diag_term * x
def _determinant(self):
return math_ops.reduce_prod(self._diag, axis=[-1])
def _log_abs_determinant(self):
log_det = math_ops.reduce_sum(
math_ops.log(math_ops.abs(self._diag)), axis=[-1])
if self.dtype.is_complex:
log_det = math_ops.cast(log_det, dtype=self.dtype)
return log_det
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
diag_term = math_ops.conj(self._diag) if adjoint else self._diag
rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
inv_diag_mat = array_ops.expand_dims(1. / diag_term, -1)
return rhs * inv_diag_mat
def _to_dense(self):
return array_ops.matrix_diag(self._diag)
def _diag_part(self):
return self.diag
def _add_to_tensor(self, x):
x_diag = array_ops.matrix_diag_part(x)
new_diag = self._diag + x_diag
return array_ops.matrix_set_diag(x, new_diag)
def _eigvals(self):
return tensor_conversion.convert_to_tensor_v2_with_dispatch(self.diag)
def _cond(self):
abs_diag = math_ops.abs(self.diag)
return (math_ops.reduce_max(abs_diag, axis=-1) /
math_ops.reduce_min(abs_diag, axis=-1))
@property
def _composite_tensor_fields(self):
return ("diag",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"diag": 1}
tensorflow/python/ops/linalg/linear_operator_full_matrix.py
+207
View File
@@ -0,0 +1,207 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` that wraps a [batch] matrix."""
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorFullMatrix"]
@tf_export("linalg.LinearOperatorFullMatrix")
@linear_operator.make_composite_tensor
class LinearOperatorFullMatrix(linear_operator.LinearOperator):
"""`LinearOperator` that wraps a [batch] matrix.
This operator wraps a [batch] matrix `A` (which is a `Tensor`) with shape
`[B1,...,Bb, M, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `M x N` matrix.
```python
# Create a 2 x 2 linear operator.
matrix = [[1., 2.], [3., 4.]]
operator = LinearOperatorFullMatrix(matrix)
operator.to_dense()
==> [[1., 2.]
[3., 4.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor
# Create a [2, 3] batch of 4 x 4 linear operators.
matrix = tf.random.normal(shape=[2, 3, 4, 4])
operator = LinearOperatorFullMatrix(matrix)
```
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [M, N], with b >= 0
x.shape = [B1,...,Bb] + [N, R], with R >= 0.
```
#### Performance
`LinearOperatorFullMatrix` has exactly the same performance as would be
achieved by using standard `TensorFlow` matrix ops. Intelligent choices are
made based on the following initialization hints.
* If `dtype` is real, and `is_self_adjoint` and `is_positive_definite`, a
Cholesky factorization is used for the determinant and solve.
In all cases, suppose `operator` is a `LinearOperatorFullMatrix` of shape
`[M, N]`, and `x.shape = [N, R]`. Then
* `operator.matmul(x)` is `O(M * N * R)`.
* If `M=N`, `operator.solve(x)` is `O(N^3 * R)`.
* If `M=N`, `operator.determinant()` is `O(N^3)`.
If instead `operator` and `x` have shape `[B1,...,Bb, M, N]` and
`[B1,...,Bb, N, R]`, every operation increases in complexity by `B1*...*Bb`.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
matrix,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorFullMatrix"):
r"""Initialize a `LinearOperatorFullMatrix`.
Args:
matrix: Shape `[B1,...,Bb, M, N]` with `b >= 0`, `M, N >= 0`.
Allowed dtypes: `float16`, `float32`, `float64`, `complex64`,
`complex128`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`.
Raises:
TypeError: If `diag.dtype` is not an allowed type.
"""
parameters = dict(
matrix=matrix,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
with ops.name_scope(name, values=[matrix]):
self._matrix = linear_operator_util.convert_nonref_to_tensor(
matrix, name="matrix")
self._check_matrix(self._matrix)
super(LinearOperatorFullMatrix, self).__init__(
dtype=self._matrix.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _check_matrix(self, matrix):
"""Static check of the `matrix` argument."""
allowed_dtypes = [
dtypes.float16,
dtypes.float32,
dtypes.float64,
dtypes.complex64,
dtypes.complex128,
]
matrix = tensor_conversion.convert_to_tensor_v2_with_dispatch(
matrix, name="matrix"
)
dtype = matrix.dtype
if dtype not in allowed_dtypes:
raise TypeError(f"Argument `matrix` must have dtype in {allowed_dtypes}. "
f"Received: {dtype}.")
if matrix.shape.ndims is not None and matrix.shape.ndims < 2:
raise ValueError(f"Argument `matrix` must have at least 2 dimensions. "
f"Received: {matrix}.")
@property
def matrix(self):
"""The matrix defining this operator."""
return self._matrix
def _shape(self):
return self._matrix.shape
def _shape_tensor(self):
return array_ops.shape(self._matrix)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
return math_ops.matmul(
self._matrix, x, adjoint_a=adjoint, adjoint_b=adjoint_arg)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
return self._dense_solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _to_dense(self):
return self._matrix
@property
def _composite_tensor_fields(self):
return ("matrix",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"matrix": 2}
tensorflow/python/ops/linalg/linear_operator_householder.py
+285
View File
@@ -0,0 +1,285 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a Householder transformation."""
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorHouseholder",]
@tf_export("linalg.LinearOperatorHouseholder")
@linear_operator.make_composite_tensor
class LinearOperatorHouseholder(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] of Householder transformations.
This operator acts like a [batch] of householder reflections with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x N` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
`LinearOperatorHouseholder` is initialized with a (batch) vector.
A Householder reflection, defined via a vector `v`, which reflects points
in `R^n` about the hyperplane orthogonal to `v` and through the origin.
```python
# Create a 2 x 2 householder transform.
vec = [1 / np.sqrt(2), 1. / np.sqrt(2)]
operator = LinearOperatorHouseholder(vec)
operator.to_dense()
==> [[0., -1.]
[-1., -0.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor
```
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
reflection_axis,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorHouseholder"):
r"""Initialize a `LinearOperatorHouseholder`.
Args:
reflection_axis: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`.
The vector defining the hyperplane to reflect about.
Allowed dtypes: `float16`, `float32`, `float64`, `complex64`,
`complex128`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. This is autoset to true
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
This is autoset to false.
is_square: Expect that this operator acts like square [batch] matrices.
This is autoset to true.
name: A name for this `LinearOperator`.
Raises:
ValueError: `is_self_adjoint` is not `True`, `is_positive_definite` is
not `False` or `is_square` is not `True`.
"""
parameters = dict(
reflection_axis=reflection_axis,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
with ops.name_scope(name, values=[reflection_axis]):
self._reflection_axis = linear_operator_util.convert_nonref_to_tensor(
reflection_axis, name="reflection_axis")
self._check_reflection_axis(self._reflection_axis)
# Check and auto-set hints.
if is_self_adjoint is False: # pylint:disable=g-bool-id-comparison
raise ValueError("A Householder operator is always self adjoint.")
else:
is_self_adjoint = True
if is_positive_definite is True: # pylint:disable=g-bool-id-comparison
raise ValueError(
"A Householder operator is always non-positive definite.")
else:
is_positive_definite = False
if is_square is False: # pylint:disable=g-bool-id-comparison
raise ValueError("A Householder operator is always square.")
is_square = True
super(LinearOperatorHouseholder, self).__init__(
dtype=self._reflection_axis.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _check_reflection_axis(self, reflection_axis):
"""Static check of reflection_axis."""
if (reflection_axis.shape.ndims is not None and
reflection_axis.shape.ndims < 1):
raise ValueError(
"Argument reflection_axis must have at least 1 dimension. "
"Found: %s" % reflection_axis)
def _shape(self):
# If d_shape = [5, 3], we return [5, 3, 3].
d_shape = self._reflection_axis.shape
return d_shape.concatenate(d_shape[-1:])
def _shape_tensor(self):
d_shape = array_ops.shape(self._reflection_axis)
k = d_shape[-1]
return array_ops.concat((d_shape, [k]), 0)
def _assert_non_singular(self):
return control_flow_ops.no_op("assert_non_singular")
def _assert_positive_definite(self):
raise errors.InvalidArgumentError(
node_def=None, op=None, message="Householder operators are always "
"non-positive definite.")
def _assert_self_adjoint(self):
return control_flow_ops.no_op("assert_self_adjoint")
def _linop_adjoint(self) -> "LinearOperatorHouseholder":
return self
def _linop_inverse(self) -> "LinearOperatorHouseholder":
return self
def _matmul(self, x, adjoint=False, adjoint_arg=False):
# Given a vector `v`, we would like to reflect `x` about the hyperplane
# orthogonal to `v` going through the origin. We first project `x` to `v`
# to get v * dot(v, x) / dot(v, v). After we project, we can reflect the
# projection about the hyperplane by flipping sign to get
# -v * dot(v, x) / dot(v, v). Finally, we can add back the component
# that is orthogonal to v. This is invariant under reflection, since the
# whole hyperplane is invariant. This component is equal to x - v * dot(v,
# x) / dot(v, v), giving the formula x - 2 * v * dot(v, x) / dot(v, v)
# for the reflection.
# Note that because this is a reflection, it lies in O(n) (for real vector
# spaces) or U(n) (for complex vector spaces), and thus is its own adjoint.
reflection_axis = tensor_conversion.convert_to_tensor_v2_with_dispatch(
self.reflection_axis
)
x = linalg.adjoint(x) if adjoint_arg else x
normalized_axis = nn.l2_normalize(reflection_axis, axis=-1)
mat = normalized_axis[..., array_ops.newaxis]
x_dot_normalized_v = math_ops.matmul(mat, x, adjoint_a=True)
return x - 2 * mat * x_dot_normalized_v
def _trace(self):
# We have (n - 1) +1 eigenvalues and a single -1 eigenvalue.
shape = self.shape_tensor()
return math_ops.cast(
self._domain_dimension_tensor(shape=shape) - 2,
self.dtype) * array_ops.ones(
shape=self._batch_shape_tensor(shape=shape), dtype=self.dtype)
def _determinant(self):
# For householder transformations, the determinant is -1.
return -array_ops.ones(shape=self.batch_shape_tensor(), dtype=self.dtype) # pylint: disable=invalid-unary-operand-type
def _log_abs_determinant(self):
# Orthogonal matrix -> log|Q| = 0.
return array_ops.zeros(shape=self.batch_shape_tensor(), dtype=self.dtype)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
# A householder reflection is a reflection, hence is idempotent. Thus we
# can just apply a matmul.
return self._matmul(rhs, adjoint, adjoint_arg)
def _to_dense(self):
reflection_axis = tensor_conversion.convert_to_tensor_v2_with_dispatch(
self.reflection_axis
)
normalized_axis = nn.l2_normalize(reflection_axis, axis=-1)
mat = normalized_axis[..., array_ops.newaxis]
matrix = -2 * math_ops.matmul(mat, mat, adjoint_b=True)
return array_ops.matrix_set_diag(
matrix, 1. + array_ops.matrix_diag_part(matrix))
def _diag_part(self):
reflection_axis = tensor_conversion.convert_to_tensor_v2_with_dispatch(
self.reflection_axis
)
normalized_axis = nn.l2_normalize(reflection_axis, axis=-1)
return 1. - 2 * normalized_axis * math_ops.conj(normalized_axis)
def _eigvals(self):
# We have (n - 1) +1 eigenvalues and a single -1 eigenvalue.
result_shape = array_ops.shape(self.reflection_axis)
n = result_shape[-1]
ones_shape = array_ops.concat([result_shape[:-1], [n - 1]], axis=-1)
neg_shape = array_ops.concat([result_shape[:-1], [1]], axis=-1)
eigvals = array_ops.ones(shape=ones_shape, dtype=self.dtype)
eigvals = array_ops.concat(
[-array_ops.ones(shape=neg_shape, dtype=self.dtype), eigvals], axis=-1) # pylint: disable=invalid-unary-operand-type
return eigvals
def _cond(self):
# Householder matrices are rotations which have condition number 1.
return array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)
@property
def reflection_axis(self):
return self._reflection_axis
@property
def _composite_tensor_fields(self):
return ("reflection_axis",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"reflection_axis": 1}
tensorflow/python/ops/linalg/linear_operator_identity.py
+933
View File
@@ -0,0 +1,933 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like the identity matrix."""
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_diag
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.ops.linalg import property_hint_util
from tensorflow.python.util.tf_export import tf_export
__all__ = [
"LinearOperatorIdentity",
"LinearOperatorScaledIdentity",
]
class BaseLinearOperatorIdentity(linear_operator.LinearOperator):
"""Base class for Identity operators."""
def _check_num_rows_possibly_add_asserts(self):
"""Static check of init arg `num_rows`, possibly add asserts."""
# Possibly add asserts.
if self._assert_proper_shapes:
self._num_rows = control_flow_ops.with_dependencies([
check_ops.assert_rank(
self._num_rows,
0,
message="Argument num_rows must be a 0-D Tensor."),
check_ops.assert_non_negative(
self._num_rows,
message="Argument num_rows must be non-negative."),
], self._num_rows)
# Static checks.
if not self._num_rows.dtype.is_integer:
raise TypeError("Argument num_rows must be integer type. Found:"
" %s" % self._num_rows)
num_rows_static = self._num_rows_static
if num_rows_static is None:
return # Cannot do any other static checks.
if num_rows_static.ndim != 0:
raise ValueError("Argument num_rows must be a 0-D Tensor. Found:"
" %s" % num_rows_static)
if num_rows_static < 0:
raise ValueError("Argument num_rows must be non-negative. Found:"
" %s" % num_rows_static)
def _min_matrix_dim(self):
"""Minimum of domain/range dimension, if statically available, else None."""
domain_dim = tensor_shape.dimension_value(self.domain_dimension)
range_dim = tensor_shape.dimension_value(self.range_dimension)
if domain_dim is None or range_dim is None:
return None
return min(domain_dim, range_dim)
def _min_matrix_dim_tensor(self):
"""Minimum of domain/range dimension, as a tensor."""
return math_ops.reduce_min(self.shape_tensor()[-2:])
def _ones_diag(self):
"""Returns the diagonal of this operator as all ones."""
if self.shape.is_fully_defined():
d_shape = self.batch_shape.concatenate([self._min_matrix_dim()])
else:
d_shape = array_ops.concat(
[self.batch_shape_tensor(),
[self._min_matrix_dim_tensor()]], axis=0)
return array_ops.ones(shape=d_shape, dtype=self.dtype)
@tf_export("linalg.LinearOperatorIdentity")
@linear_operator.make_composite_tensor
class LinearOperatorIdentity(BaseLinearOperatorIdentity):
"""`LinearOperator` acting like a [batch] square identity matrix.
This operator acts like a [batch] identity matrix `A` with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x N` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
`LinearOperatorIdentity` is initialized with `num_rows`, and optionally
`batch_shape`, and `dtype` arguments. If `batch_shape` is `None`, this
operator efficiently passes through all arguments. If `batch_shape` is
provided, broadcasting may occur, which will require making copies.
```python
# Create a 2 x 2 identity matrix.
operator = LinearOperatorIdentity(num_rows=2, dtype=tf.float32)
operator.to_dense()
==> [[1., 0.]
[0., 1.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> 0.
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor, same as x.
y = tf.random.normal(shape=[3, 2, 4])
# Note that y.shape is compatible with operator.shape because operator.shape
# is broadcast to [3, 2, 2].
# This broadcast does NOT require copying data, since we can infer that y
# will be passed through without changing shape. We are always able to infer
# this if the operator has no batch_shape.
x = operator.solve(y)
==> Shape [3, 2, 4] Tensor, same as y.
# Create a 2-batch of 2x2 identity matrices
operator = LinearOperatorIdentity(num_rows=2, batch_shape=[2])
operator.to_dense()
==> [[[1., 0.]
[0., 1.]],
[[1., 0.]
[0., 1.]]]
# Here, even though the operator has a batch shape, the input is the same as
# the output, so x can be passed through without a copy. The operator is able
# to detect that no broadcast is necessary because both x and the operator
# have statically defined shape.
x = ... Shape [2, 2, 3]
operator.matmul(x)
==> Shape [2, 2, 3] Tensor, same as x
# Here the operator and x have different batch_shape, and are broadcast.
# This requires a copy, since the output is different size than the input.
x = ... Shape [1, 2, 3]
operator.matmul(x)
==> Shape [2, 2, 3] Tensor, equal to [x, x]
```
### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
### Performance
If `batch_shape` initialization arg is `None`:
* `operator.matmul(x)` is `O(1)`
* `operator.solve(x)` is `O(1)`
* `operator.determinant()` is `O(1)`
If `batch_shape` initialization arg is provided, and static checks cannot
rule out the need to broadcast:
* `operator.matmul(x)` is `O(D1*...*Dd*N*R)`
* `operator.solve(x)` is `O(D1*...*Dd*N*R)`
* `operator.determinant()` is `O(B1*...*Bb)`
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
num_rows,
batch_shape=None,
dtype=None,
is_non_singular=True,
is_self_adjoint=True,
is_positive_definite=True,
is_square=True,
assert_proper_shapes=False,
name="LinearOperatorIdentity"):
r"""Initialize a `LinearOperatorIdentity`.
The `LinearOperatorIdentity` is initialized with arguments defining `dtype`
and shape.
This operator is able to broadcast the leading (batch) dimensions, which
sometimes requires copying data. If `batch_shape` is `None`, the operator
can take arguments of any batch shape without copying. See examples.
Args:
num_rows: Scalar non-negative integer `Tensor`. Number of rows in the
corresponding identity matrix.
batch_shape: Optional `1-D` integer `Tensor`. The shape of the leading
dimensions. If `None`, this operator has no leading dimensions.
dtype: Data type of the matrix that this operator represents.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
assert_proper_shapes: Python `bool`. If `False`, only perform static
checks that initialization and method arguments have proper shape.
If `True`, and static checks are inconclusive, add asserts to the graph.
name: A name for this `LinearOperator`
Raises:
ValueError: If `num_rows` is determined statically to be non-scalar, or
negative.
ValueError: If `batch_shape` is determined statically to not be 1-D, or
negative.
ValueError: If any of the following is not `True`:
`{is_self_adjoint, is_non_singular, is_positive_definite}`.
TypeError: If `num_rows` or `batch_shape` is ref-type (e.g. Variable).
"""
parameters = dict(
num_rows=num_rows,
batch_shape=batch_shape,
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
assert_proper_shapes=assert_proper_shapes,
name=name)
dtype = dtype or dtypes.float32
self._assert_proper_shapes = assert_proper_shapes
with ops.name_scope(name):
dtype = dtypes.as_dtype(dtype)
if not is_self_adjoint:
raise ValueError("An identity operator is always self adjoint.")
if not is_non_singular:
raise ValueError("An identity operator is always non-singular.")
if not is_positive_definite:
raise ValueError("An identity operator is always positive-definite.")
if not is_square:
raise ValueError("An identity operator is always square.")
super(LinearOperatorIdentity, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
linear_operator_util.assert_not_ref_type(num_rows, "num_rows")
linear_operator_util.assert_not_ref_type(batch_shape, "batch_shape")
self._num_rows = linear_operator_util.shape_tensor(
num_rows, name="num_rows")
self._num_rows_static = tensor_util.constant_value(self._num_rows)
self._check_num_rows_possibly_add_asserts()
if batch_shape is None:
self._batch_shape_arg = None
else:
self._batch_shape_arg = linear_operator_util.shape_tensor(
batch_shape, name="batch_shape_arg")
self._batch_shape_static = tensor_util.constant_value(
self._batch_shape_arg)
self._check_batch_shape_possibly_add_asserts()
def _shape(self):
matrix_shape = tensor_shape.TensorShape((self._num_rows_static,
self._num_rows_static))
if self._batch_shape_arg is None:
return matrix_shape
batch_shape = tensor_shape.TensorShape(self._batch_shape_static)
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
matrix_shape = array_ops_stack.stack(
(self._num_rows, self._num_rows), axis=0)
if self._batch_shape_arg is None:
return matrix_shape
return array_ops.concat((self._batch_shape_arg, matrix_shape), 0)
def _linop_adjoint(self) -> "LinearOperatorIdentity":
return self
def _linop_cholesky(self) -> "LinearOperatorIdentity":
return LinearOperatorIdentity(
num_rows=self._num_rows, # pylint: disable=protected-access
batch_shape=self.batch_shape,
dtype=self.dtype,
is_non_singular=True,
is_self_adjoint=True,
is_positive_definite=True,
is_square=True)
def _linop_inverse(self) -> "LinearOperatorIdentity":
return self
def _linop_matmul(
self,
left_operator: "LinearOperatorIdentity",
right_operator: linear_operator.LinearOperator,
) -> "LinearOperatorIdentity":
del left_operator
return right_operator
def _linop_solve(
self,
left_operator: "LinearOperatorIdentity",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
del left_operator
return right_operator
def _assert_non_singular(self):
return control_flow_ops.no_op("assert_non_singular")
def _assert_positive_definite(self):
return control_flow_ops.no_op("assert_positive_definite")
def _assert_self_adjoint(self):
return control_flow_ops.no_op("assert_self_adjoint")
def _possibly_broadcast_batch_shape(self, x):
"""Return 'x', possibly after broadcasting the leading dimensions."""
# If we have no batch shape, our batch shape broadcasts with everything!
if self._batch_shape_arg is None:
return x
# Static attempt:
# If we determine that no broadcast is necessary, pass x through
# If we need a broadcast, add to an array of zeros.
#
# special_shape is the shape that, when broadcast with x's shape, will give
# the correct broadcast_shape. Note that
# We have already verified the second to last dimension of self.shape
# matches x's shape in assert_compatible_matrix_dimensions.
# Also, the final dimension of 'x' can have any shape.
# Therefore, the final two dimensions of special_shape are 1's.
special_shape = self.batch_shape.concatenate([1, 1])
bshape = array_ops.broadcast_static_shape(x.shape, special_shape)
if special_shape.is_fully_defined():
# bshape.is_fully_defined iff special_shape.is_fully_defined.
if bshape == x.shape:
return x
# Use the built in broadcasting of addition.
zeros = array_ops.zeros(shape=special_shape, dtype=self.dtype)
return x + zeros
# Dynamic broadcast:
# Always add to an array of zeros, rather than using a "cond", since a
# cond would require copying data from GPU --> CPU.
special_shape = array_ops.concat((self.batch_shape_tensor(), [1, 1]), 0)
zeros = array_ops.zeros(shape=special_shape, dtype=self.dtype)
return x + zeros
def _matmul(self, x, adjoint=False, adjoint_arg=False):
# Note that adjoint has no effect since this matrix is self-adjoint.
x = linalg.adjoint(x) if adjoint_arg else x
if self._assert_proper_shapes:
aps = linear_operator_util.assert_compatible_matrix_dimensions(self, x)
x = control_flow_ops.with_dependencies([aps], x)
return self._possibly_broadcast_batch_shape(x)
def _determinant(self):
return array_ops.ones(shape=self.batch_shape_tensor(), dtype=self.dtype)
def _log_abs_determinant(self):
return array_ops.zeros(shape=self.batch_shape_tensor(), dtype=self.dtype)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
return self._matmul(rhs, adjoint_arg=adjoint_arg)
def _trace(self):
# Get Tensor of all ones of same shape as self.batch_shape.
if self.batch_shape.is_fully_defined():
batch_of_ones = array_ops.ones(shape=self.batch_shape, dtype=self.dtype)
else:
batch_of_ones = array_ops.ones(
shape=self.batch_shape_tensor(), dtype=self.dtype)
if self._min_matrix_dim() is not None:
return self._min_matrix_dim() * batch_of_ones
else:
return (math_ops.cast(self._min_matrix_dim_tensor(), self.dtype) *
batch_of_ones)
def _diag_part(self):
return self._ones_diag()
def add_to_tensor(self, mat, name="add_to_tensor"):
"""Add matrix represented by this operator to `mat`. Equiv to `I + mat`.
Args:
mat: `Tensor` with same `dtype` and shape broadcastable to `self`.
name: A name to give this `Op`.
Returns:
A `Tensor` with broadcast shape and same `dtype` as `self`.
"""
with self._name_scope(name): # pylint: disable=not-callable
mat = tensor_conversion.convert_to_tensor_v2_with_dispatch(
mat, name="mat"
)
mat_diag = array_ops.matrix_diag_part(mat)
new_diag = 1 + mat_diag
return array_ops.matrix_set_diag(mat, new_diag)
def _eigvals(self):
return self._ones_diag()
def _cond(self):
return array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)
def _check_num_rows_possibly_add_asserts(self):
"""Static check of init arg `num_rows`, possibly add asserts."""
# Possibly add asserts.
if self._assert_proper_shapes:
self._num_rows = control_flow_ops.with_dependencies([
check_ops.assert_rank(
self._num_rows,
0,
message="Argument num_rows must be a 0-D Tensor."),
check_ops.assert_non_negative(
self._num_rows,
message="Argument num_rows must be non-negative."),
], self._num_rows)
# Static checks.
if not self._num_rows.dtype.is_integer:
raise TypeError("Argument num_rows must be integer type. Found:"
" %s" % self._num_rows)
num_rows_static = self._num_rows_static
if num_rows_static is None:
return # Cannot do any other static checks.
if num_rows_static.ndim != 0:
raise ValueError("Argument num_rows must be a 0-D Tensor. Found:"
" %s" % num_rows_static)
if num_rows_static < 0:
raise ValueError("Argument num_rows must be non-negative. Found:"
" %s" % num_rows_static)
def _check_batch_shape_possibly_add_asserts(self):
"""Static check of init arg `batch_shape`, possibly add asserts."""
if self._batch_shape_arg is None:
return
# Possibly add asserts
if self._assert_proper_shapes:
self._batch_shape_arg = control_flow_ops.with_dependencies([
check_ops.assert_rank(
self._batch_shape_arg,
1,
message="Argument batch_shape must be a 1-D Tensor."),
check_ops.assert_non_negative(
self._batch_shape_arg,
message="Argument batch_shape must be non-negative."),
], self._batch_shape_arg)
# Static checks
if not self._batch_shape_arg.dtype.is_integer:
raise TypeError("Argument batch_shape must be integer type. Found:"
" %s" % self._batch_shape_arg)
if self._batch_shape_static is None:
return # Cannot do any other static checks.
if self._batch_shape_static.ndim != 1:
raise ValueError("Argument batch_shape must be a 1-D Tensor. Found:"
" %s" % self._batch_shape_static)
if np.any(self._batch_shape_static < 0):
raise ValueError("Argument batch_shape must be non-negative. Found:"
"%s" % self._batch_shape_static)
@property
def _composite_tensor_prefer_static_fields(self):
return ("num_rows", "batch_shape")
@property
def _composite_tensor_fields(self):
return ("num_rows", "batch_shape", "dtype", "assert_proper_shapes")
def __getitem__(self, slices):
# Slice the batch shape and return a new LinearOperatorIdentity.
# Use a proxy shape and slice it. Use this as the new batch shape
new_batch_shape = array_ops.shape(
array_ops.ones(self._batch_shape_arg)[slices])
parameters = dict(self.parameters, batch_shape=new_batch_shape)
return LinearOperatorIdentity(**parameters)
@tf_export("linalg.LinearOperatorScaledIdentity")
@linear_operator.make_composite_tensor
class LinearOperatorScaledIdentity(BaseLinearOperatorIdentity):
"""`LinearOperator` acting like a scaled [batch] identity matrix `A = c I`.
This operator acts like a scaled [batch] identity matrix `A` with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
a scaled version of the `N x N` identity matrix.
`LinearOperatorIdentity` is initialized with `num_rows`, and a `multiplier`
(a `Tensor`) of shape `[B1,...,Bb]`. `N` is set to `num_rows`, and the
`multiplier` determines the scale for each batch member.
```python
# Create a 2 x 2 scaled identity matrix.
operator = LinearOperatorIdentity(num_rows=2, multiplier=3.)
operator.to_dense()
==> [[3., 0.]
[0., 3.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> 2 * Log[3]
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> 3 * x
y = tf.random.normal(shape=[3, 2, 4])
# Note that y.shape is compatible with operator.shape because operator.shape
# is broadcast to [3, 2, 2].
x = operator.solve(y)
==> 3 * x
# Create a 2-batch of 2x2 identity matrices
operator = LinearOperatorIdentity(num_rows=2, multiplier=5.)
operator.to_dense()
==> [[[5., 0.]
[0., 5.]],
[[5., 0.]
[0., 5.]]]
x = ... Shape [2, 2, 3]
operator.matmul(x)
==> 5 * x
# Here the operator and x have different batch_shape, and are broadcast.
x = ... Shape [1, 2, 3]
operator.matmul(x)
==> 5 * x
```
### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
### Performance
* `operator.matmul(x)` is `O(D1*...*Dd*N*R)`
* `operator.solve(x)` is `O(D1*...*Dd*N*R)`
* `operator.determinant()` is `O(D1*...*Dd)`
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
num_rows,
multiplier,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=True,
assert_proper_shapes=False,
name="LinearOperatorScaledIdentity"):
r"""Initialize a `LinearOperatorScaledIdentity`.
The `LinearOperatorScaledIdentity` is initialized with `num_rows`, which
determines the size of each identity matrix, and a `multiplier`,
which defines `dtype`, batch shape, and scale of each matrix.
This operator is able to broadcast the leading (batch) dimensions.
Args:
num_rows: Scalar non-negative integer `Tensor`. Number of rows in the
corresponding identity matrix.
multiplier: `Tensor` of shape `[B1,...,Bb]`, or `[]` (a scalar).
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
assert_proper_shapes: Python `bool`. If `False`, only perform static
checks that initialization and method arguments have proper shape. If
`True`, and static checks are inconclusive, add asserts to the graph.
name: A name for this `LinearOperator`
Raises:
ValueError: If `num_rows` is determined statically to be non-scalar, or
negative.
ValueError: If the dtype of `multiplier` is complex and `is_self_adjoint`
is `False`.
ValueError: If `is_positive_definite` is `True` and `is_non_singular`
is `False`.
"""
parameters = dict(
num_rows=num_rows,
multiplier=multiplier,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
assert_proper_shapes=assert_proper_shapes,
name=name)
self._assert_proper_shapes = assert_proper_shapes
with ops.name_scope(name, values=[multiplier, num_rows]):
self._multiplier = linear_operator_util.convert_nonref_to_tensor(
multiplier, name="multiplier")
# Check and auto-set hints.
if not self._multiplier.dtype.is_complex:
if is_self_adjoint is False: # pylint: disable=g-bool-id-comparison
raise ValueError("A real diagonal operator is always self adjoint.")
else:
is_self_adjoint = True
if not is_square:
raise ValueError("A ScaledIdentity operator is always square.")
linear_operator_util.assert_not_ref_type(num_rows, "num_rows")
super(LinearOperatorScaledIdentity, self).__init__(
dtype=self._multiplier.dtype.base_dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
self._num_rows = linear_operator_util.shape_tensor(
num_rows, name="num_rows")
self._num_rows_static = tensor_util.constant_value(self._num_rows)
self._check_num_rows_possibly_add_asserts()
self._num_rows_cast_to_dtype = math_ops.cast(self._num_rows, self.dtype)
self._num_rows_cast_to_real_dtype = math_ops.cast(self._num_rows,
self.dtype.real_dtype)
def _shape(self):
matrix_shape = tensor_shape.TensorShape((self._num_rows_static,
self._num_rows_static))
batch_shape = self.multiplier.shape
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
matrix_shape = array_ops_stack.stack(
(self._num_rows, self._num_rows), axis=0)
batch_shape = array_ops.shape(self.multiplier)
return array_ops.concat((batch_shape, matrix_shape), 0)
def _assert_non_singular(self):
return check_ops.assert_positive(
math_ops.abs(self.multiplier), message="LinearOperator was singular")
def _assert_positive_definite(self):
return check_ops.assert_positive(
math_ops.real(self.multiplier),
message="LinearOperator was not positive definite.")
def _assert_self_adjoint(self):
imag_multiplier = math_ops.imag(self.multiplier)
return check_ops.assert_equal(
array_ops.zeros_like(imag_multiplier),
imag_multiplier,
message="LinearOperator was not self-adjoint")
def _make_multiplier_matrix(self, conjugate=False):
# Shape [B1,...Bb, 1, 1]
multiplier_matrix = array_ops.expand_dims(
array_ops.expand_dims(self.multiplier, -1), -1)
if conjugate:
multiplier_matrix = math_ops.conj(multiplier_matrix)
return multiplier_matrix
def _linop_adjoint(self) -> "LinearOperatorScaledIdentity":
multiplier = self.multiplier
if multiplier.dtype.is_complex:
multiplier = math_ops.conj(multiplier)
return LinearOperatorScaledIdentity(
num_rows=self._num_rows,
multiplier=multiplier,
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_cholesky(self) -> "LinearOperatorScaledIdentity":
return LinearOperatorScaledIdentity(
num_rows=self._num_rows,
multiplier=math_ops.sqrt(self.multiplier),
is_non_singular=True,
is_self_adjoint=True,
is_positive_definite=True,
is_square=True)
def _linop_inverse(self) -> "LinearOperatorScaledIdentity":
return LinearOperatorScaledIdentity(
num_rows=self._num_rows,
multiplier=1. / self.multiplier,
is_non_singular=self.is_non_singular,
is_self_adjoint=True,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_matmul(
self,
left_operator: "LinearOperatorScaledIdentity",
right_operator: linear_operator.LinearOperator,
) -> "LinearOperatorScaledIdentity":
is_non_singular = property_hint_util.combined_non_singular_hint(
left_operator, right_operator)
is_self_adjoint = property_hint_util.combined_commuting_self_adjoint_hint(
left_operator, right_operator)
is_positive_definite = (
property_hint_util.combined_commuting_positive_definite_hint(
left_operator, right_operator))
if isinstance(right_operator, LinearOperatorScaledIdentity):
return LinearOperatorScaledIdentity(
num_rows=left_operator.domain_dimension_tensor(),
multiplier=left_operator.multiplier * right_operator.multiplier,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
elif isinstance(right_operator, linear_operator_diag.LinearOperatorDiag):
return linear_operator_diag.LinearOperatorDiag(
diag=right_operator.diag * left_operator.multiplier,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
else:
return super()._linop_matmul(left_operator, right_operator)
def _linop_solve(
self,
left_operator: "LinearOperatorScaledIdentity",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
is_non_singular = property_hint_util.combined_non_singular_hint(
left_operator, right_operator)
is_self_adjoint = property_hint_util.combined_commuting_self_adjoint_hint(
left_operator, right_operator)
is_positive_definite = (
property_hint_util.combined_commuting_positive_definite_hint(
left_operator, right_operator))
if isinstance(right_operator, LinearOperatorScaledIdentity):
return LinearOperatorScaledIdentity(
num_rows=left_operator.domain_dimension_tensor(),
multiplier=right_operator.multiplier / left_operator.multiplier,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
elif isinstance(right_operator, linear_operator_diag.LinearOperatorDiag):
return linear_operator_diag.LinearOperatorDiag(
diag=right_operator.diag / left_operator.multiplier,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=True)
else:
return super()._linop_solve(left_operator, right_operator)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
x = linalg.adjoint(x) if adjoint_arg else x
if self._assert_proper_shapes:
aps = linear_operator_util.assert_compatible_matrix_dimensions(self, x)
x = control_flow_ops.with_dependencies([aps], x)
return x * self._make_multiplier_matrix(conjugate=adjoint)
def _determinant(self):
return self.multiplier**self._num_rows_cast_to_dtype
def _log_abs_determinant(self):
return self._num_rows_cast_to_real_dtype * math_ops.log(
math_ops.abs(self.multiplier))
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
if self._assert_proper_shapes:
aps = linear_operator_util.assert_compatible_matrix_dimensions(self, rhs)
rhs = control_flow_ops.with_dependencies([aps], rhs)
return rhs / self._make_multiplier_matrix(conjugate=adjoint)
def _trace(self):
# Get Tensor of all ones of same shape as self.batch_shape.
if self.batch_shape.is_fully_defined():
batch_of_ones = array_ops.ones(shape=self.batch_shape, dtype=self.dtype)
else:
batch_of_ones = array_ops.ones(
shape=self.batch_shape_tensor(), dtype=self.dtype)
if self._min_matrix_dim() is not None:
return self.multiplier * self._min_matrix_dim() * batch_of_ones
else:
return (self.multiplier * math_ops.cast(self._min_matrix_dim_tensor(),
self.dtype) * batch_of_ones)
def _diag_part(self):
return self._ones_diag() * self.multiplier[..., array_ops.newaxis]
def add_to_tensor(self, mat, name="add_to_tensor"):
"""Add matrix represented by this operator to `mat`. Equiv to `I + mat`.
Args:
mat: `Tensor` with same `dtype` and shape broadcastable to `self`.
name: A name to give this `Op`.
Returns:
A `Tensor` with broadcast shape and same `dtype` as `self`.
"""
with self._name_scope(name): # pylint: disable=not-callable
# Shape [B1,...,Bb, 1]
multiplier_vector = array_ops.expand_dims(self.multiplier, -1)
# Shape [C1,...,Cc, M, M]
mat = tensor_conversion.convert_to_tensor_v2_with_dispatch(
mat, name="mat"
)
# Shape [C1,...,Cc, M]
mat_diag = array_ops.matrix_diag_part(mat)
# multiplier_vector broadcasts here.
new_diag = multiplier_vector + mat_diag
return array_ops.matrix_set_diag(mat, new_diag)
def _eigvals(self):
return self._ones_diag() * self.multiplier[..., array_ops.newaxis]
def _cond(self):
# Condition number for a scalar time identity matrix is one, except when the
# scalar is zero.
return array_ops.where_v2(
math_ops.equal(self._multiplier, 0.),
math_ops.cast(np.nan, dtype=self.dtype),
math_ops.cast(1., dtype=self.dtype))
@property
def multiplier(self):
"""The [batch] scalar `Tensor`, `c` in `cI`."""
return self._multiplier
@property
def _composite_tensor_prefer_static_fields(self):
return ("num_rows",)
@property
def _composite_tensor_fields(self):
return ("num_rows", "multiplier", "assert_proper_shapes")
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"multiplier": 0}
tensorflow/python/ops/linalg/linear_operator_inversion.py
+231
View File
@@ -0,0 +1,231 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Inverts a non-singular `LinearOperator`."""
from tensorflow.python.framework import ops
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorInversion"]
@tf_export("linalg.LinearOperatorInversion")
@linear_operator.make_composite_tensor
class LinearOperatorInversion(linear_operator.LinearOperator):
"""`LinearOperator` representing the inverse of another operator.
This operator represents the inverse of another operator.
```python
# Create a 2 x 2 linear operator.
operator = LinearOperatorFullMatrix([[1., 0.], [0., 2.]])
operator_inv = LinearOperatorInversion(operator)
operator_inv.to_dense()
==> [[1., 0.]
[0., 0.5]]
operator_inv.shape
==> [2, 2]
operator_inv.log_abs_determinant()
==> - log(2)
x = ... Shape [2, 4] Tensor
operator_inv.matmul(x)
==> Shape [2, 4] Tensor, equal to operator.solve(x)
```
#### Performance
The performance of `LinearOperatorInversion` depends on the underlying
operators performance: `solve` and `matmul` are swapped, and determinant is
inverted.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
operator,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name=None):
r"""Initialize a `LinearOperatorInversion`.
`LinearOperatorInversion` is initialized with an operator `A`. The `solve`
and `matmul` methods are effectively swapped. E.g.
```
A = MyLinearOperator(...)
B = LinearOperatorInversion(A)
x = [....] # a vector
assert A.matvec(x) == B.solvevec(x)
```
Args:
operator: `LinearOperator` object. If `operator.is_non_singular == False`,
an exception is raised. We do allow `operator.is_non_singular == None`,
in which case this operator will have `is_non_singular == None`.
Similarly for `is_self_adjoint` and `is_positive_definite`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`. Default is `operator.name +
"_inv"`.
Raises:
ValueError: If `operator.is_non_singular` is False.
"""
parameters = dict(
operator=operator,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
self._operator = operator
# Auto-set and check hints.
if operator.is_non_singular is False or is_non_singular is False:
raise ValueError(
f"Argument `is_non_singular` or argument `operator` must have "
f"supplied hint `is_non_singular` equal to `True` or `None`. "
f"Found `operator.is_non_singular`: {operator.is_non_singular}, "
f"`is_non_singular`: {is_non_singular}.")
if operator.is_square is False or is_square is False:
raise ValueError(
f"Argument `is_square` or argument `operator` must have supplied "
f"hint `is_square` equal to `True` or `None`. Found "
f"`operator.is_square`: {operator.is_square}, "
f"`is_square`: {is_square}.")
# The congruency of is_non_singular and is_self_adjoint was checked in the
# base operator. Other hints are, in this special case of inversion, ones
# that must be the same for base/derived operator.
combine_hint = (
linear_operator_util.use_operator_or_provided_hint_unless_contradicting)
is_square = combine_hint(
operator, "is_square", is_square,
"An operator is square if and only if its inverse is square.")
is_non_singular = combine_hint(
operator, "is_non_singular", is_non_singular,
"An operator is non-singular if and only if its inverse is "
"non-singular.")
is_self_adjoint = combine_hint(
operator, "is_self_adjoint", is_self_adjoint,
"An operator is self-adjoint if and only if its inverse is "
"self-adjoint.")
is_positive_definite = combine_hint(
operator, "is_positive_definite", is_positive_definite,
"An operator is positive-definite if and only if its inverse is "
"positive-definite.")
# Initialization.
if name is None:
name = operator.name + "_inv"
with ops.name_scope(name):
super(LinearOperatorInversion, self).__init__(
dtype=operator.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
@property
def operator(self) -> "LinearOperatorInversion":
"""The operator before inversion."""
return self._operator
def _linop_inverse(self) -> linear_operator.LinearOperator:
return self.operator
def _linop_solve(
self,
left_operator: "LinearOperatorInversion",
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
"""Solve inverse of generic `LinearOperator`s."""
return left_operator.operator.matmul(right_operator)
def _assert_non_singular(self):
return self.operator.assert_non_singular()
def _assert_positive_definite(self):
return self.operator.assert_positive_definite()
def _assert_self_adjoint(self):
return self.operator.assert_self_adjoint()
def _shape(self):
return self.operator.shape
def _shape_tensor(self):
return self.operator.shape_tensor()
def _matmul(self, x, adjoint=False, adjoint_arg=False):
return self.operator.solve(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _determinant(self):
return 1. / self.operator.determinant()
def _log_abs_determinant(self):
return -1. * self.operator.log_abs_determinant()
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
return self.operator.matmul(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
def _eigvals(self):
return 1. / self.operator.eigvals()
def _cond(self):
return self.operator.cond()
@property
def _composite_tensor_fields(self):
return ("operator",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"operator": 0}
tensorflow/python/ops/linalg/linear_operator_kronecker.py
+538
View File
@@ -0,0 +1,538 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Construct the Kronecker product of one or more `LinearOperators`."""
from tensorflow.python.framework import common_shapes
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorKronecker"]
def _prefer_static_shape(x):
if x.shape.is_fully_defined():
return x.shape
return array_ops.shape(x)
def _prefer_static_concat_shape(first_shape, second_shape_int_list):
"""Concatenate a shape with a list of integers as statically as possible.
Args:
first_shape: `TensorShape` or `Tensor` instance. If a `TensorShape`,
`first_shape.is_fully_defined()` must return `True`.
second_shape_int_list: `list` of scalar integer `Tensor`s.
Returns:
`Tensor` representing concatenating `first_shape` and
`second_shape_int_list` as statically as possible.
"""
second_shape_int_list_static = [
tensor_util.constant_value(s) for s in second_shape_int_list]
if (isinstance(first_shape, tensor_shape.TensorShape) and
all(s is not None for s in second_shape_int_list_static)):
return first_shape.concatenate(second_shape_int_list_static)
return array_ops.concat([first_shape, second_shape_int_list], axis=0)
@tf_export("linalg.LinearOperatorKronecker")
@linear_operator.make_composite_tensor
class LinearOperatorKronecker(linear_operator.LinearOperator):
"""Kronecker product between two `LinearOperators`.
This operator composes one or more linear operators `[op1,...,opJ]`,
building a new `LinearOperator` representing the Kronecker product:
`op1 x op2 x .. opJ` (we omit parentheses as the Kronecker product is
associative).
If `opj` has shape `batch_shape_j + [M_j, N_j]`, then the composed operator
will have shape equal to `broadcast_batch_shape + [prod M_j, prod N_j]`,
where the product is over all operators.
```python
# Create a 4 x 4 linear operator composed of two 2 x 2 operators.
operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]])
operator_2 = LinearOperatorFullMatrix([[1., 0.], [2., 1.]])
operator = LinearOperatorKronecker([operator_1, operator_2])
operator.to_dense()
==> [[1., 0., 2., 0.],
[2., 1., 4., 2.],
[3., 0., 4., 0.],
[6., 3., 8., 4.]]
operator.shape
==> [4, 4]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [4, 2] Tensor
operator.matmul(x)
==> Shape [4, 2] Tensor
# Create a [2, 3] batch of 4 x 5 linear operators.
matrix_45 = tf.random.normal(shape=[2, 3, 4, 5])
operator_45 = LinearOperatorFullMatrix(matrix)
# Create a [2, 3] batch of 5 x 6 linear operators.
matrix_56 = tf.random.normal(shape=[2, 3, 5, 6])
operator_56 = LinearOperatorFullMatrix(matrix_56)
# Compose to create a [2, 3] batch of 20 x 30 operators.
operator_large = LinearOperatorKronecker([operator_45, operator_56])
# Create a shape [2, 3, 20, 2] vector.
x = tf.random.normal(shape=[2, 3, 6, 2])
operator_large.matmul(x)
==> Shape [2, 3, 30, 2] Tensor
```
#### Performance
The performance of `LinearOperatorKronecker` on any operation is equal to
the sum of the individual operators' operations.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
operators,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name=None):
r"""Initialize a `LinearOperatorKronecker`.
`LinearOperatorKronecker` is initialized with a list of operators
`[op_1,...,op_J]`.
Args:
operators: Iterable of `LinearOperator` objects, each with
the same `dtype` and composable shape, representing the Kronecker
factors.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix\
#Extension_for_non_symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`. Default is the individual
operators names joined with `_x_`.
Raises:
TypeError: If all operators do not have the same `dtype`.
ValueError: If `operators` is empty.
"""
parameters = dict(
operators=operators,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
# Validate operators.
check_ops.assert_proper_iterable(operators)
operators = list(operators)
if not operators:
raise ValueError(f"Argument `operators` must be a list of >=1 operators. "
f"Received: {operators}.")
self._operators = operators
# Validate dtype.
dtype = operators[0].dtype
for operator in operators:
if operator.dtype != dtype:
name_type = (str((o.name, o.dtype)) for o in operators)
raise TypeError(
f"Expected every operation in argument `operators` to have the "
f"same dtype. Received {list(name_type)}.")
# Auto-set and check hints.
# A Kronecker product is invertible, if and only if all factors are
# invertible.
if all(operator.is_non_singular for operator in operators):
if is_non_singular is False:
raise ValueError(
f"The Kronecker product of non-singular operators is always "
f"non-singular. Expected argument `is_non_singular` to be True. "
f"Received: {is_non_singular}.")
is_non_singular = True
if all(operator.is_self_adjoint for operator in operators):
if is_self_adjoint is False:
raise ValueError(
f"The Kronecker product of self-adjoint operators is always "
f"self-adjoint. Expected argument `is_self_adjoint` to be True. "
f"Received: {is_self_adjoint}.")
is_self_adjoint = True
# The eigenvalues of a Kronecker product are equal to the products of eigen
# values of the corresponding factors.
if all(operator.is_positive_definite for operator in operators):
if is_positive_definite is False:
raise ValueError(
f"The Kronecker product of positive-definite operators is always "
f"positive-definite. Expected argument `is_positive_definite` to "
f"be True. Received: {is_positive_definite}.")
is_positive_definite = True
if name is None:
name = operators[0].name
for operator in operators[1:]:
name += "_x_" + operator.name
with ops.name_scope(name):
super(LinearOperatorKronecker, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
@property
def operators(self):
return self._operators
def _shape(self):
# Get final matrix shape.
domain_dimension = self.operators[0].domain_dimension
for operator in self.operators[1:]:
domain_dimension = domain_dimension * operator.domain_dimension
range_dimension = self.operators[0].range_dimension
for operator in self.operators[1:]:
range_dimension = range_dimension * operator.range_dimension
matrix_shape = tensor_shape.TensorShape([
range_dimension, domain_dimension])
# Get broadcast batch shape.
# broadcast_shape checks for compatibility.
batch_shape = self.operators[0].batch_shape
for operator in self.operators[1:]:
batch_shape = common_shapes.broadcast_shape(
batch_shape, operator.batch_shape)
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
domain_dimension = self.operators[0].domain_dimension_tensor()
for operator in self.operators[1:]:
domain_dimension = domain_dimension * operator.domain_dimension_tensor()
range_dimension = self.operators[0].range_dimension_tensor()
for operator in self.operators[1:]:
range_dimension = range_dimension * operator.range_dimension_tensor()
matrix_shape = [range_dimension, domain_dimension]
# Get broadcast batch shape.
# broadcast_shape checks for compatibility.
batch_shape = self.operators[0].batch_shape_tensor()
for operator in self.operators[1:]:
batch_shape = array_ops.broadcast_dynamic_shape(
batch_shape, operator.batch_shape_tensor())
return array_ops.concat((batch_shape, matrix_shape), 0)
def _linop_adjoint(self) -> "LinearOperatorKronecker":
return LinearOperatorKronecker(
operators=[operator.adjoint() for operator in self.operators],
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _linop_cholesky(self) -> "LinearOperatorKronecker":
# Cholesky decomposition of a Kronecker product is the Kronecker product
# of cholesky decompositions.
return LinearOperatorKronecker(
operators=[operator.cholesky() for operator in self.operators],
is_non_singular=True,
is_self_adjoint=None, # Let the operators passed in decide.
is_square=True)
def _linop_inverse(self) -> "LinearOperatorKronecker":
# Inverse decomposition of a Kronecker product is the Kronecker product
# of inverse decompositions.
return LinearOperatorKronecker(
operators=[
operator.inverse() for operator in self.operators],
is_non_singular=self.is_non_singular,
is_self_adjoint=self.is_self_adjoint,
is_positive_definite=self.is_positive_definite,
is_square=True)
def _solve_matmul_internal(
self,
x,
solve_matmul_fn,
adjoint=False,
adjoint_arg=False):
# We heavily rely on Roth's column Lemma [1]:
# (A x B) * vec X = vec BXA^T
# where vec stacks all the columns of the matrix under each other.
# In our case, we use a variant of the lemma that is row-major
# friendly: (A x B) * vec' X = vec' AXB^T
# Where vec' reshapes a matrix into a vector. We can repeatedly apply this
# for a collection of kronecker products.
# Given that (A x B)^-1 = A^-1 x B^-1 and (A x B)^T = A^T x B^T, we can
# use the above to compute multiplications, solves with any composition of
# transposes.
output = x
if adjoint_arg:
if self.dtype.is_complex:
output = math_ops.conj(output)
else:
output = linalg.transpose(output)
for o in reversed(self.operators):
# Statically compute the reshape.
if adjoint:
operator_dimension = o.range_dimension_tensor()
else:
operator_dimension = o.domain_dimension_tensor()
output_shape = _prefer_static_shape(output)
if tensor_util.constant_value(operator_dimension) is not None:
operator_dimension = tensor_util.constant_value(operator_dimension)
if output.shape[-2] is not None and output.shape[-1] is not None:
dim = int(output.shape[-2] * output_shape[-1] // operator_dimension)
else:
dim = math_ops.cast(
output_shape[-2] * output_shape[-1] // operator_dimension,
dtype=dtypes.int32)
output_shape = _prefer_static_concat_shape(
output_shape[:-2], [dim, operator_dimension])
output = array_ops.reshape(output, shape=output_shape)
# Conjugate because we are trying to compute A @ B^T, but
# `LinearOperator` only supports `adjoint_arg`.
if self.dtype.is_complex:
output = math_ops.conj(output)
output = solve_matmul_fn(
o, output, adjoint=adjoint, adjoint_arg=True)
if adjoint_arg:
col_dim = _prefer_static_shape(x)[-2]
else:
col_dim = _prefer_static_shape(x)[-1]
if adjoint:
row_dim = self.domain_dimension_tensor()
else:
row_dim = self.range_dimension_tensor()
matrix_shape = [row_dim, col_dim]
output = array_ops.reshape(
output,
_prefer_static_concat_shape(
_prefer_static_shape(output)[:-2], matrix_shape))
if x.shape.is_fully_defined():
if adjoint_arg:
column_dim = x.shape[-2]
else:
column_dim = x.shape[-1]
broadcast_batch_shape = common_shapes.broadcast_shape(
x.shape[:-2], self.batch_shape)
if adjoint:
matrix_dimensions = [self.domain_dimension, column_dim]
else:
matrix_dimensions = [self.range_dimension, column_dim]
output.set_shape(broadcast_batch_shape.concatenate(
matrix_dimensions))
return output
def _matmul(self, x, adjoint=False, adjoint_arg=False):
def matmul_fn(o, x, adjoint, adjoint_arg):
return o.matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
return self._solve_matmul_internal(
x=x,
solve_matmul_fn=matmul_fn,
adjoint=adjoint,
adjoint_arg=adjoint_arg)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
def solve_fn(o, rhs, adjoint, adjoint_arg):
return o.solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
return self._solve_matmul_internal(
x=rhs,
solve_matmul_fn=solve_fn,
adjoint=adjoint,
adjoint_arg=adjoint_arg)
def _determinant(self):
# Note that we have |X1 x X2| = |X1| ** n * |X2| ** m, where X1 is an m x m
# matrix, and X2 is an n x n matrix. We can iteratively apply this property
# to get the determinant of |X1 x X2 x X3 ...|. If T is the product of the
# domain dimension of all operators, then we have:
# |X1 x X2 x X3 ...| =
# |X1| ** (T / m) * |X2 x X3 ... | ** m =
# |X1| ** (T / m) * |X2| ** (m * (T / m) / n) * ... =
# |X1| ** (T / m) * |X2| ** (T / n) * | X3 x X4... | ** (m * n)
# And by doing induction we have product(|X_i| ** (T / dim(X_i))).
total = self.domain_dimension_tensor()
determinant = 1.
for operator in self.operators:
determinant = determinant * operator.determinant() ** math_ops.cast(
total / operator.domain_dimension_tensor(),
dtype=operator.dtype)
return determinant
def _log_abs_determinant(self):
# This will be sum((total / dim(x_i)) * log |X_i|)
total = self.domain_dimension_tensor()
log_abs_det = 0.
for operator in self.operators:
log_abs_det += operator.log_abs_determinant() * math_ops.cast(
total / operator.domain_dimension_tensor(),
dtype=operator.dtype)
return log_abs_det
def _trace(self):
# tr(A x B) = tr(A) * tr(B)
trace = 1.
for operator in self.operators:
trace = trace * operator.trace()
return trace
def _diag_part(self):
diag_part = self.operators[0].diag_part()
for operator in self.operators[1:]:
diag_part = diag_part[..., :, array_ops.newaxis]
op_diag_part = operator.diag_part()[..., array_ops.newaxis, :]
diag_part = diag_part * op_diag_part
diag_part = array_ops.reshape(
diag_part,
shape=array_ops.concat(
[array_ops.shape(diag_part)[:-2], [-1]], axis=0))
if self.range_dimension > self.domain_dimension:
diag_dimension = self.domain_dimension
else:
diag_dimension = self.range_dimension
diag_part.set_shape(
self.batch_shape.concatenate(diag_dimension))
return diag_part
def _to_dense(self):
product = self.operators[0].to_dense()
for operator in self.operators[1:]:
# Product has shape [B, R1, 1, C1, 1].
product = product[
..., :, array_ops.newaxis, :, array_ops.newaxis]
# Operator has shape [B, 1, R2, 1, C2].
op_to_mul = operator.to_dense()[
..., array_ops.newaxis, :, array_ops.newaxis, :]
# This is now [B, R1, R2, C1, C2].
product = product * op_to_mul
# Now merge together dimensions to get [B, R1 * R2, C1 * C2].
product_shape = _prefer_static_shape(product)
shape = _prefer_static_concat_shape(
product_shape[:-4],
[product_shape[-4] * product_shape[-3],
product_shape[-2] * product_shape[-1]])
product = array_ops.reshape(product, shape=shape)
product.set_shape(self.shape)
return product
def _eigvals(self):
# This will be the kronecker product of all the eigenvalues.
# Note: It doesn't matter which kronecker product it is, since every
# kronecker product of the same matrices are similar.
eigvals = [operator.eigvals() for operator in self.operators]
# Now compute the kronecker product
product = eigvals[0]
for eigval in eigvals[1:]:
# Product has shape [B, R1, 1].
product = product[..., array_ops.newaxis]
# Eigval has shape [B, 1, R2]. Produces shape [B, R1, R2].
product = product * eigval[..., array_ops.newaxis, :]
# Reshape to [B, R1 * R2]
product = array_ops.reshape(
product,
shape=array_ops.concat([array_ops.shape(product)[:-2], [-1]], axis=0))
product.set_shape(self.shape[:-1])
return product
def _assert_non_singular(self):
if all(operator.is_square for operator in self.operators):
asserts = [operator.assert_non_singular() for operator in self.operators]
return control_flow_ops.group(asserts)
else:
raise errors.InvalidArgumentError(
node_def=None,
op=None,
message="All Kronecker factors must be square for the product to be "
"invertible. Expected hint `is_square` to be True for every operator "
"in argument `operators`.")
def _assert_self_adjoint(self):
if all(operator.is_square for operator in self.operators):
asserts = [operator.assert_self_adjoint() for operator in self.operators]
return control_flow_ops.group(asserts)
else:
raise errors.InvalidArgumentError(
node_def=None,
op=None,
message="All Kronecker factors must be square for the product to be "
"invertible. Expected hint `is_square` to be True for every operator "
"in argument `operators`.")
@property
def _composite_tensor_fields(self):
return ("operators",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"operators": [0] * len(self.operators)}
tensorflow/python/ops/linalg/linear_operator_low_rank_update.py
+511
View File
@@ -0,0 +1,511 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Perturb a `LinearOperator` with a rank `K` update."""
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_diag
from tensorflow.python.ops.linalg import linear_operator_identity
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util.tf_export import tf_export
__all__ = [
"LinearOperatorLowRankUpdate",
]
@tf_export("linalg.LinearOperatorLowRankUpdate")
@linear_operator.make_composite_tensor
class LinearOperatorLowRankUpdate(linear_operator.LinearOperator):
"""Perturb a `LinearOperator` with a rank `K` update.
This operator acts like a [batch] matrix `A` with shape
`[B1,...,Bb, M, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `M x N` matrix.
`LinearOperatorLowRankUpdate` represents `A = L + U D V^H`, where
```
L, is a LinearOperator representing [batch] M x N matrices
U, is a [batch] M x K matrix. Typically K << M.
D, is a [batch] K x K matrix.
V, is a [batch] N x K matrix. Typically K << N.
V^H is the Hermitian transpose (adjoint) of V.
```
If `M = N`, determinants and solves are done using the matrix determinant
lemma and Woodbury identities, and thus require L and D to be non-singular.
Solves and determinants will be attempted unless the "is_non_singular"
property of L and D is False.
In the event that L and D are positive-definite, and U = V, solves and
determinants can be done using a Cholesky factorization.
```python
# Create a 3 x 3 diagonal linear operator.
diag_operator = LinearOperatorDiag(
diag_update=[1., 2., 3.], is_non_singular=True, is_self_adjoint=True,
is_positive_definite=True)
# Perturb with a rank 2 perturbation
operator = LinearOperatorLowRankUpdate(
operator=diag_operator,
u=[[1., 2.], [-1., 3.], [0., 0.]],
diag_update=[11., 12.],
v=[[1., 2.], [-1., 3.], [10., 10.]])
operator.shape
==> [3, 3]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [3, 4] Tensor
operator.matmul(x)
==> Shape [3, 4] Tensor
```
### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [M, N], with b >= 0
x.shape = [B1,...,Bb] + [N, R], with R >= 0.
```
### Performance
Suppose `operator` is a `LinearOperatorLowRankUpdate` of shape `[M, N]`,
made from a rank `K` update of `base_operator` which performs `.matmul(x)` on
`x` having `x.shape = [N, R]` with `O(L_matmul*N*R)` complexity (and similarly
for `solve`, `determinant`. Then, if `x.shape = [N, R]`,
* `operator.matmul(x)` is `O(L_matmul*N*R + K*N*R)`
and if `M = N`,
* `operator.solve(x)` is `O(L_matmul*N*R + N*K*R + K^2*R + K^3)`
* `operator.determinant()` is `O(L_determinant + L_solve*N*K + K^2*N + K^3)`
If instead `operator` and `x` have shape `[B1,...,Bb, M, N]` and
`[B1,...,Bb, N, R]`, every operation increases in complexity by `B1*...*Bb`.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular`, `self_adjoint`, `positive_definite`,
`diag_update_positive` and `square`. These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
base_operator,
u,
diag_update=None,
v=None,
is_diag_update_positive=None,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorLowRankUpdate"):
"""Initialize a `LinearOperatorLowRankUpdate`.
This creates a `LinearOperator` of the form `A = L + U D V^H`, with
`L` a `LinearOperator`, `U, V` both [batch] matrices, and `D` a [batch]
diagonal matrix.
If `L` is non-singular, solves and determinants are available.
Solves/determinants both involve a solve/determinant of a `K x K` system.
In the event that L and D are self-adjoint positive-definite, and U = V,
this can be done using a Cholesky factorization. The user should set the
`is_X` matrix property hints, which will trigger the appropriate code path.
Args:
base_operator: Shape `[B1,...,Bb, M, N]`.
u: Shape `[B1,...,Bb, M, K]` `Tensor` of same `dtype` as `base_operator`.
This is `U` above.
diag_update: Optional shape `[B1,...,Bb, K]` `Tensor` with same `dtype`
as `base_operator`. This is the diagonal of `D` above.
Defaults to `D` being the identity operator.
v: Optional `Tensor` of same `dtype` as `u` and shape `[B1,...,Bb, N, K]`
Defaults to `v = u`, in which case the perturbation is symmetric.
If `M != N`, then `v` must be set since the perturbation is not square.
is_diag_update_positive: Python `bool`.
If `True`, expect `diag_update > 0`.
is_non_singular: Expect that this operator is non-singular.
Default is `None`, unless `is_positive_definite` is auto-set to be
`True` (see below).
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. Default is `None`, unless `base_operator` is self-adjoint
and `v = None` (meaning `u=v`), in which case this defaults to `True`.
is_positive_definite: Expect that this operator is positive definite.
Default is `None`, unless `base_operator` is positive-definite
`v = None` (meaning `u=v`), and `is_diag_update_positive`, in which case
this defaults to `True`.
Note that we say an operator is positive definite when the quadratic
form `x^H A x` has positive real part for all nonzero `x`.
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`.
Raises:
ValueError: If `is_X` flags are set in an inconsistent way.
"""
parameters = dict(
base_operator=base_operator,
u=u,
diag_update=diag_update,
v=v,
is_diag_update_positive=is_diag_update_positive,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
dtype = base_operator.dtype
if diag_update is not None:
if is_diag_update_positive and dtype.is_complex:
logging.warn("Note: setting is_diag_update_positive with a complex "
"dtype means that diagonal is real and positive.")
if diag_update is None:
if is_diag_update_positive is False:
raise ValueError(
"Default diagonal is the identity, which is positive. However, "
"user set 'is_diag_update_positive' to False.")
is_diag_update_positive = True
# In this case, we can use a Cholesky decomposition to help us solve/det.
self._use_cholesky = (
base_operator.is_positive_definite and base_operator.is_self_adjoint
and is_diag_update_positive
and v is None)
# Possibly auto-set some characteristic flags from None to True.
# If the Flags were set (by the user) incorrectly to False, then raise.
if base_operator.is_self_adjoint and v is None and not dtype.is_complex:
if is_self_adjoint is False:
raise ValueError(
"A = L + UDU^H, with L self-adjoint and D real diagonal. Since"
" UDU^H is self-adjoint, this must be a self-adjoint operator.")
is_self_adjoint = True
# The condition for using a cholesky is sufficient for SPD, and
# we no weaker choice of these hints leads to SPD. Therefore,
# the following line reads "if hints indicate SPD..."
if self._use_cholesky:
if (
is_positive_definite is False
or is_self_adjoint is False
or is_non_singular is False):
raise ValueError(
"Arguments imply this is self-adjoint positive-definite operator.")
is_positive_definite = True
is_self_adjoint = True
with ops.name_scope(name):
# Create U and V.
self._u = linear_operator_util.convert_nonref_to_tensor(u, name="u")
if v is None:
self._v = self._u
else:
self._v = linear_operator_util.convert_nonref_to_tensor(v, name="v")
if diag_update is None:
self._diag_update = None
else:
self._diag_update = linear_operator_util.convert_nonref_to_tensor(
diag_update, name="diag_update")
# Create base_operator L.
self._base_operator = base_operator
super(LinearOperatorLowRankUpdate, self).__init__(
dtype=self._base_operator.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
# Create the diagonal operator D.
self._set_diag_operators(diag_update, is_diag_update_positive)
self._is_diag_update_positive = is_diag_update_positive
self._check_shapes()
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 _set_diag_operators(self, diag_update, is_diag_update_positive):
"""Set attributes self._diag_update and self._diag_operator."""
if diag_update is not None:
self._diag_operator = linear_operator_diag.LinearOperatorDiag(
self._diag_update, is_positive_definite=is_diag_update_positive)
else:
if tensor_shape.dimension_value(self.u.shape[-1]) is not None:
r = tensor_shape.dimension_value(self.u.shape[-1])
else:
r = array_ops.shape(self.u)[-1]
self._diag_operator = linear_operator_identity.LinearOperatorIdentity(
num_rows=r, dtype=self.dtype)
@property
def u(self):
"""If this operator is `A = L + U D V^H`, this is the `U`."""
return self._u
@property
def v(self):
"""If this operator is `A = L + U D V^H`, this is the `V`."""
return self._v
@property
def is_diag_update_positive(self):
"""If this operator is `A = L + U D V^H`, this hints `D > 0` elementwise."""
return self._is_diag_update_positive
@property
def diag_update(self):
"""If this operator is `A = L + U D V^H`, this is the diagonal of `D`."""
return self._diag_update
@property
def diag_operator(self):
"""If this operator is `A = L + U D V^H`, this is `D`."""
return self._diag_operator
@property
def base_operator(self):
"""If this operator is `A = L + U D V^H`, this is the `L`."""
return self._base_operator
def _assert_self_adjoint(self):
# Recall this operator is:
# A = L + UDV^H.
# So in one case self-adjoint depends only on L
if self.u is self.v and self.diag_update is None:
return self.base_operator.assert_self_adjoint()
# In all other cases, sufficient conditions for self-adjoint can be found
# efficiently. However, those conditions are not necessary conditions.
return super(LinearOperatorLowRankUpdate, self).assert_self_adjoint()
def _shape(self):
batch_shape = array_ops.broadcast_static_shape(
self.base_operator.batch_shape,
self.diag_operator.batch_shape)
batch_shape = array_ops.broadcast_static_shape(
batch_shape,
self.u.shape[:-2])
batch_shape = array_ops.broadcast_static_shape(
batch_shape,
self.v.shape[:-2])
return batch_shape.concatenate(self.base_operator.shape[-2:])
def _shape_tensor(self):
batch_shape = array_ops.broadcast_dynamic_shape(
self.base_operator.batch_shape_tensor(),
self.diag_operator.batch_shape_tensor())
batch_shape = array_ops.broadcast_dynamic_shape(
batch_shape,
array_ops.shape(self.u)[:-2])
batch_shape = array_ops.broadcast_dynamic_shape(
batch_shape,
array_ops.shape(self.v)[:-2])
return array_ops.concat(
[batch_shape, self.base_operator.shape_tensor()[-2:]], axis=0)
def _get_uv_as_tensors(self):
"""Get (self.u, self.v) as tensors (in case they were refs)."""
u = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.u)
if self.v is self.u:
v = u
else:
v = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.v)
return u, v
def _matmul(self, x, adjoint=False, adjoint_arg=False):
u, v = self._get_uv_as_tensors()
l = self.base_operator
d = self.diag_operator
leading_term = l.matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg)
if adjoint:
uh_x = math_ops.matmul(u, x, adjoint_a=True, adjoint_b=adjoint_arg)
d_uh_x = d.matmul(uh_x, adjoint=adjoint)
v_d_uh_x = math_ops.matmul(v, d_uh_x)
return leading_term + v_d_uh_x
else:
vh_x = math_ops.matmul(v, x, adjoint_a=True, adjoint_b=adjoint_arg)
d_vh_x = d.matmul(vh_x, adjoint=adjoint)
u_d_vh_x = math_ops.matmul(u, d_vh_x)
return leading_term + u_d_vh_x
def _determinant(self):
if self.is_positive_definite:
return math_ops.exp(self.log_abs_determinant())
# The matrix determinant lemma gives
# https://en.wikipedia.org/wiki/Matrix_determinant_lemma
# det(L + UDV^H) = det(D^{-1} + V^H L^{-1} U) det(D) det(L)
# = det(C) det(D) det(L)
# where C is sometimes known as the capacitance matrix,
# C := D^{-1} + V^H L^{-1} U
u, v = self._get_uv_as_tensors()
det_c = linalg_ops.matrix_determinant(self._make_capacitance(u=u, v=v))
det_d = self.diag_operator.determinant()
det_l = self.base_operator.determinant()
return det_c * det_d * det_l
def _diag_part(self):
# [U D V^T]_{ii} = sum_{jk} U_{ij} D_{jk} V_{ik}
# = sum_{j} U_{ij} D_{jj} V_{ij}
u, v = self._get_uv_as_tensors()
product = u * math_ops.conj(v)
if self.diag_update is not None:
product *= array_ops.expand_dims(self.diag_update, axis=-2)
return (
math_ops.reduce_sum(product, axis=-1) + self.base_operator.diag_part())
def _log_abs_determinant(self):
u, v = self._get_uv_as_tensors()
# Recall
# det(L + UDV^H) = det(D^{-1} + V^H L^{-1} U) det(D) det(L)
# = det(C) det(D) det(L)
log_abs_det_d = self.diag_operator.log_abs_determinant()
log_abs_det_l = self.base_operator.log_abs_determinant()
if self._use_cholesky:
chol_cap_diag = array_ops.matrix_diag_part(
linalg_ops.cholesky(self._make_capacitance(u=u, v=v)))
log_abs_det_c = 2 * math_ops.reduce_sum(
math_ops.log(chol_cap_diag), axis=[-1])
else:
det_c = linalg_ops.matrix_determinant(self._make_capacitance(u=u, v=v))
log_abs_det_c = math_ops.log(math_ops.abs(det_c))
if self.dtype.is_complex:
log_abs_det_c = math_ops.cast(log_abs_det_c, dtype=self.dtype)
return log_abs_det_c + log_abs_det_d + log_abs_det_l
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
if self.base_operator.is_non_singular is False:
raise ValueError(
"Solve not implemented unless this is a perturbation of a "
"non-singular LinearOperator.")
# The Woodbury formula gives:
# https://en.wikipedia.org/wiki/Woodbury_matrix_identity
# (L + UDV^H)^{-1}
# = L^{-1} - L^{-1} U (D^{-1} + V^H L^{-1} U)^{-1} V^H L^{-1}
# = L^{-1} - L^{-1} U C^{-1} V^H L^{-1}
# where C is the capacitance matrix, C := D^{-1} + V^H L^{-1} U
# Note also that, with ^{-H} being the inverse of the adjoint,
# (L + UDV^H)^{-H}
# = L^{-H} - L^{-H} V C^{-H} U^H L^{-H}
l = self.base_operator
if adjoint:
# If adjoint, U and V have flipped roles in the operator.
v, u = self._get_uv_as_tensors()
# Capacitance should still be computed with u=self.u and v=self.v, which
# after the "flip" on the line above means u=v, v=u. I.e. no need to
# "flip" in the capacitance call, since the call to
# matrix_solve_with_broadcast below is done with the `adjoint` argument,
# and this takes care of things.
capacitance = self._make_capacitance(u=v, v=u)
else:
u, v = self._get_uv_as_tensors()
capacitance = self._make_capacitance(u=u, v=v)
# L^{-1} rhs
linv_rhs = l.solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg)
# V^H L^{-1} rhs
vh_linv_rhs = math_ops.matmul(v, linv_rhs, adjoint_a=True)
# C^{-1} V^H L^{-1} rhs
if self._use_cholesky:
capinv_vh_linv_rhs = linalg_ops.cholesky_solve(
linalg_ops.cholesky(capacitance), vh_linv_rhs)
else:
capinv_vh_linv_rhs = linear_operator_util.matrix_solve_with_broadcast(
capacitance, vh_linv_rhs, adjoint=adjoint)
# U C^{-1} V^H M^{-1} rhs
u_capinv_vh_linv_rhs = math_ops.matmul(u, capinv_vh_linv_rhs)
# L^{-1} U C^{-1} V^H L^{-1} rhs
linv_u_capinv_vh_linv_rhs = l.solve(u_capinv_vh_linv_rhs, adjoint=adjoint)
# L^{-1} - L^{-1} U C^{-1} V^H L^{-1}
return linv_rhs - linv_u_capinv_vh_linv_rhs
def _make_capacitance(self, u, v):
# C := D^{-1} + V^H L^{-1} U
# which is sometimes known as the "capacitance" matrix.
# L^{-1} U
linv_u = self.base_operator.solve(u)
# V^H L^{-1} U
vh_linv_u = math_ops.matmul(v, linv_u, adjoint_a=True)
# D^{-1} + V^H L^{-1} V
capacitance = self._diag_operator.inverse().add_to_tensor(vh_linv_u)
return capacitance
@property
def _composite_tensor_fields(self):
return ("base_operator", "u", "diag_update", "v", "is_diag_update_positive")
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {
"base_operator": 0,
"u": 2,
"diag_update": 1,
"v": 2
}
tensorflow/python/ops/linalg/linear_operator_lower_triangular.py
+249
View File
@@ -0,0 +1,249 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a lower triangular matrix."""
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.ops.linalg import property_hint_util
from tensorflow.python.util.tf_export import tf_export
__all__ = [
"LinearOperatorLowerTriangular",
]
@tf_export("linalg.LinearOperatorLowerTriangular")
@linear_operator.make_composite_tensor
class LinearOperatorLowerTriangular(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] square lower triangular matrix.
This operator acts like a [batch] lower triangular matrix `A` with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x N` matrix.
`LinearOperatorLowerTriangular` is initialized with a `Tensor` having
dimensions `[B1,...,Bb, N, N]`. The upper triangle of the last two
dimensions is ignored.
```python
# Create a 2 x 2 lower-triangular linear operator.
tril = [[1., 2.], [3., 4.]]
operator = LinearOperatorLowerTriangular(tril)
# The upper triangle is ignored.
operator.to_dense()
==> [[1., 0.]
[3., 4.]]
operator.shape
==> [2, 2]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor
# Create a [2, 3] batch of 4 x 4 linear operators.
tril = tf.random.normal(shape=[2, 3, 4, 4])
operator = LinearOperatorLowerTriangular(tril)
```
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [B1,...,Bb] + [N, R], with R >= 0.
```
#### Performance
Suppose `operator` is a `LinearOperatorLowerTriangular` of shape `[N, N]`,
and `x.shape = [N, R]`. Then
* `operator.matmul(x)` involves `N^2 * R` multiplications.
* `operator.solve(x)` involves `N * R` size `N` back-substitutions.
* `operator.determinant()` involves a size `N` `reduce_prod`.
If instead `operator` and `x` have shape `[B1,...,Bb, N, N]` and
`[B1,...,Bb, N, R]`, every operation increases in complexity by `B1*...*Bb`.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
tril,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorLowerTriangular"):
r"""Initialize a `LinearOperatorLowerTriangular`.
Args:
tril: Shape `[B1,...,Bb, N, N]` with `b >= 0`, `N >= 0`.
The lower triangular part of `tril` defines this operator. The strictly
upper triangle is ignored.
is_non_singular: Expect that this operator is non-singular.
This operator is non-singular if and only if its diagonal elements are
all non-zero.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. This operator is self-adjoint only if it is diagonal with
real-valued diagonal entries. In this case it is advised to use
`LinearOperatorDiag`.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`.
Raises:
ValueError: If `is_square` is `False`.
"""
parameters = dict(
tril=tril,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
if is_square is False:
raise ValueError(
"Only square lower triangular operators supported at this time.")
is_square = True
with ops.name_scope(name, values=[tril]):
self._tril = linear_operator_util.convert_nonref_to_tensor(tril,
name="tril")
self._check_tril(self._tril)
super(LinearOperatorLowerTriangular, self).__init__(
dtype=self._tril.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
@property
def tril(self):
"""The lower triangular matrix defining this operator."""
return self._tril
def _check_tril(self, tril):
"""Static check of the `tril` argument."""
if tril.shape.ndims is not None and tril.shape.ndims < 2:
raise ValueError(
"Argument tril must have at least 2 dimensions. Found: %s"
% tril)
def _get_tril(self):
"""Gets the `tril` kwarg, with upper part zero-d out."""
return array_ops.matrix_band_part(self._tril, -1, 0)
def _get_diag(self):
"""Gets the diagonal part of `tril` kwarg."""
return array_ops.matrix_diag_part(self._tril)
def _shape(self):
return self._tril.shape
def _shape_tensor(self):
return array_ops.shape(self._tril)
def _assert_non_singular(self):
return linear_operator_util.assert_no_entries_with_modulus_zero(
self._get_diag(),
message="Singular operator: Diagonal contained zero values.")
def _matmul(self, x, adjoint=False, adjoint_arg=False):
return math_ops.matmul(
self._get_tril(), x, adjoint_a=adjoint, adjoint_b=adjoint_arg)
def _linop_matmul(
self,
left_operator: linear_operator.LinearOperator,
right_operator: linear_operator.LinearOperator,
) -> linear_operator.LinearOperator:
# instance check of linear_operator_diag.LinearOperatorDiag
if (hasattr(right_operator, "_check_diag") and
isinstance(left_operator, LinearOperatorLowerTriangular)):
is_non_singular = property_hint_util.combined_non_singular_hint(
left_operator, right_operator)
is_self_adjoint = property_hint_util.combined_commuting_self_adjoint_hint(
left_operator, right_operator)
return LinearOperatorLowerTriangular(
tril=left_operator.to_dense() * array_ops.expand_dims(
right_operator.diag, axis=-2),
is_non_singular=is_non_singular,
# This is safe to do since the Triangular matrix is only self-adjoint
# when it is a diagonal matrix, and hence commutes.
is_self_adjoint=is_self_adjoint,
is_positive_definite=None,
is_square=True)
return super()._linop_matmul(left_operator, right_operator)
def _determinant(self):
return math_ops.reduce_prod(self._get_diag(), axis=[-1])
def _log_abs_determinant(self):
return math_ops.reduce_sum(
math_ops.log(math_ops.abs(self._get_diag())), axis=[-1])
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
return linalg.triangular_solve(
self._get_tril(), rhs, lower=True, adjoint=adjoint)
def _to_dense(self):
return self._get_tril()
def _eigvals(self):
return self._get_diag()
@property
def _composite_tensor_fields(self):
return ("tril",)
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"tril": 2}
tensorflow/python/ops/linalg/linear_operator_permutation.py
+271
View File
@@ -0,0 +1,271 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a permutation matrix."""
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import sort_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorPermutation",]
@tf_export("linalg.LinearOperatorPermutation")
@linear_operator.make_composite_tensor
class LinearOperatorPermutation(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] of permutation matrices.
This operator acts like a [batch] of permutations with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x N` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
`LinearOperatorPermutation` is initialized with a (batch) vector.
A permutation, is defined by an integer vector `v` whose values are unique
and are in the range `[0, ... n]`. Applying the permutation on an input
matrix has the folllowing meaning: the value of `v` at index `i`
says to move the `v[i]`-th row of the input matrix to the `i`-th row.
Because all values are unique, this will result in a permutation of the
rows the input matrix. Note, that the permutation vector `v` has the same
semantics as `tf.transpose`.
```python
# Create a 3 x 3 permutation matrix that swaps the last two columns.
vec = [0, 2, 1]
operator = LinearOperatorPermutation(vec)
operator.to_dense()
==> [[1., 0., 0.]
[0., 0., 1.]
[0., 1., 0.]]
operator.shape
==> [3, 3]
# This will be zero.
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [3, 4] Tensor
operator.matmul(x)
==> Shape [3, 4] Tensor
```
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
perm,
dtype=dtypes.float32,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorPermutation"):
r"""Initialize a `LinearOperatorPermutation`.
Args:
perm: Shape `[B1,...,Bb, N]` Integer `Tensor` with `b >= 0`
`N >= 0`. An integer vector that represents the permutation to apply.
Note that this argument is same as `tf.transpose`. However, this
permutation is applied on the rows, while the permutation in
`tf.transpose` is applied on the dimensions of the `Tensor`. `perm`
is required to have unique entries from `{0, 1, ... N-1}`.
dtype: The `dtype` of arguments to this operator. Default: `float32`.
Allowed dtypes: `float16`, `float32`, `float64`, `complex64`,
`complex128`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. This is autoset to true
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
This is autoset to false.
is_square: Expect that this operator acts like square [batch] matrices.
This is autoset to true.
name: A name for this `LinearOperator`.
Raises:
ValueError: `is_self_adjoint` is not `True`, `is_positive_definite` is
not `False` or `is_square` is not `True`.
"""
parameters = dict(
perm=perm,
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
with ops.name_scope(name, values=[perm]):
self._perm = linear_operator_util.convert_nonref_to_tensor(
perm, name="perm")
self._check_perm(self._perm)
# Check and auto-set hints.
if is_non_singular is False: # pylint:disable=g-bool-id-comparison
raise ValueError(f"A Permutation operator is always non-singular. "
f"Expected argument `is_non_singular` to be True. "
f"Received: {is_non_singular}.")
if is_square is False: # pylint:disable=g-bool-id-comparison
raise ValueError(f"A Permutation operator is always square. "
f"Expected argument `is_square` to be True. "
f"Received: {is_square}.")
is_square = True
super(LinearOperatorPermutation, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _check_perm(self, perm):
"""Static check of perm."""
if (perm.shape.ndims is not None and perm.shape.ndims < 1):
raise ValueError(f"Argument `perm` must have at least 1 dimension. "
f"Received: {perm}.")
if not perm.dtype.is_integer:
raise TypeError(f"Argument `perm` must be integer dtype. "
f"Received: {perm}.")
# Check that the permutation satisfies the uniqueness constraint.
static_perm = tensor_util.constant_value(perm)
if static_perm is not None:
sorted_perm = np.sort(static_perm, axis=-1)
if np.any(sorted_perm != np.arange(0, static_perm.shape[-1])):
raise ValueError(
f"Argument `perm` must be a vector of unique integers from "
f"0 to {static_perm.shape[-1] - 1}.")
def _shape(self):
perm_shape = self._perm.shape
return perm_shape.concatenate(perm_shape[-1:])
def _shape_tensor(self):
perm_shape = array_ops.shape(self._perm)
k = perm_shape[-1]
return array_ops.concat((perm_shape, [k]), 0)
def _assert_non_singular(self):
return control_flow_ops.no_op("assert_non_singular")
def _domain_dimension_tensor(self, perm=None):
perm = perm if perm is not None else self.perm
return array_ops.shape(perm)[-1]
def _matmul(self, x, adjoint=False, adjoint_arg=False):
perm = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.perm)
if adjoint and not self.is_self_adjoint:
# TODO(srvasude): invert_permutation doesn't work on batches so we use
# argsort.
perm = sort_ops.argsort(perm, axis=-1)
x = linalg.adjoint(x) if adjoint_arg else x
# We need to broadcast x and the permutation since tf.gather doesn't
# broadcast.
broadcast_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(x)[:-1], array_ops.shape(perm))
k = array_ops.shape(x)[-1]
broadcast_x_shape = array_ops.concat([broadcast_shape, [k]], axis=-1)
x = array_ops.broadcast_to(x, broadcast_x_shape)
perm = array_ops.broadcast_to(perm, broadcast_shape)
m = array_ops.shape(x)[-2]
x = array_ops.reshape(x, [-1, m, k])
perm = array_ops.reshape(perm, [-1, m])
y = array_ops.gather(x, perm, axis=-2, batch_dims=1)
return array_ops.reshape(y, broadcast_x_shape)
# TODO(srvasude): Permutation parity is equivalent to the determinant.
def _log_abs_determinant(self):
# Permutation matrices have determinant +/- 1.
return array_ops.zeros(shape=self.batch_shape_tensor(), dtype=self.dtype)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
# The inverse of a permutation matrix is the transpose matrix.
# Apply a matmul and flip the adjoint bit.
return self._matmul(rhs, adjoint=(not adjoint), adjoint_arg=adjoint_arg)
def _to_dense(self):
perm = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.perm)
return math_ops.cast(math_ops.equal(
math_ops.range(0, self._domain_dimension_tensor(perm)),
perm[..., array_ops.newaxis]), self.dtype)
def _diag_part(self):
perm = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.perm)
return math_ops.cast(math_ops.equal(
math_ops.range(0, self._domain_dimension_tensor(perm)),
perm), self.dtype)
def _cond(self):
# Permutation matrices are rotations which have condition number 1.
return array_ops.ones(self.batch_shape_tensor(), dtype=self.dtype)
@property
def perm(self):
return self._perm
@property
def _composite_tensor_fields(self):
return ("perm", "dtype")
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"perm": 1}
tensorflow/python/ops/linalg/linear_operator_test_util.py
+1436
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/linalg/linear_operator_toeplitz.py
+299
View File
@@ -0,0 +1,299 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a Toeplitz matrix."""
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_circulant
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.ops.signal import fft_ops
from tensorflow.python.util.tf_export import tf_export
__all__ = ["LinearOperatorToeplitz",]
@tf_export("linalg.LinearOperatorToeplitz")
@linear_operator.make_composite_tensor
class LinearOperatorToeplitz(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] of toeplitz matrices.
This operator acts like a [batch] Toeplitz matrix `A` with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x N` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
#### Description in terms of toeplitz matrices
Toeplitz means that `A` has constant diagonals. Hence, `A` can be generated
with two vectors. One represents the first column of the matrix, and the
other represents the first row.
Below is a 4 x 4 example:
```
A = |a b c d|
|e a b c|
|f e a b|
|g f e a|
```
#### Example of a Toeplitz operator.
```python
# Create a 3 x 3 Toeplitz operator.
col = [1., 2., 3.]
row = [1., 4., -9.]
operator = LinearOperatorToeplitz(col, row)
operator.to_dense()
==> [[1., 4., -9.],
[2., 1., 4.],
[3., 2., 1.]]
operator.shape
==> [3, 3]
operator.log_abs_determinant()
==> scalar Tensor
x = ... Shape [3, 4] Tensor
operator.matmul(x)
==> Shape [3, 4] Tensor
```
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
col,
row,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name="LinearOperatorToeplitz"):
r"""Initialize a `LinearOperatorToeplitz`.
Args:
col: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`. The first
column of the operator. Allowed dtypes: `float16`, `float32`, `float64`,
`complex64`, `complex128`. Note that the first entry of `col` is assumed
to be the same as the first entry of `row`.
row: Shape `[B1,...,Bb, N]` `Tensor` with `b >= 0` `N >= 0`. The first row
of the operator. Allowed dtypes: `float16`, `float32`, `float64`,
`complex64`, `complex128`. Note that the first entry of `row` is assumed
to be the same as the first entry of `col`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. If `diag.dtype` is real, this is auto-set to `True`.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
Should not be `False` for that only square Toeplitz operators currently
supported.
name: A name for this `LinearOperator`.
Raises:
ValueError: If `is_square` is `False`.
ValueError: If `is_positive_definite` is `True` and
`is_non_singular` is `False`.
"""
parameters = dict(
col=col,
row=row,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
with ops.name_scope(name, values=[row, col]):
self._row = linear_operator_util.convert_nonref_to_tensor(row, name="row")
self._col = linear_operator_util.convert_nonref_to_tensor(col, name="col")
self._check_row_col(self._row, self._col)
if is_square is False: # pylint:disable=g-bool-id-comparison
raise ValueError("Only square Toeplitz operators currently supported.")
is_square = True
super(LinearOperatorToeplitz, self).__init__(
dtype=self._row.dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _check_row_col(self, row, col):
"""Static check of row and column."""
for name, tensor in [["row", row], ["col", col]]:
if tensor.shape.ndims is not None and tensor.shape.ndims < 1:
raise ValueError("Argument {} must have at least 1 dimension. "
"Found: {}".format(name, tensor))
if row.shape[-1] is not None and col.shape[-1] is not None:
if row.shape[-1] != col.shape[-1]:
raise ValueError(
"Expected square matrix, got row and col with mismatched "
"dimensions.")
def _shape(self):
# If d_shape = [5, 3], we return [5, 3, 3].
v_shape = array_ops.broadcast_static_shape(
self.row.shape, self.col.shape)
return v_shape.concatenate(v_shape[-1:])
def _shape_tensor(self, row=None, col=None):
row = self.row if row is None else row
col = self.col if col is None else col
v_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(row),
array_ops.shape(col))
k = v_shape[-1]
return array_ops.concat((v_shape, [k]), 0)
def _assert_self_adjoint(self):
return check_ops.assert_equal(
self.row,
self.col,
message=("row and col are not the same, and "
"so this operator is not self-adjoint."))
# TODO(srvasude): Add efficient solver and determinant calculations to this
# class (based on Levinson recursion.)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
# Given a Toeplitz matrix, we can embed it in a Circulant matrix to perform
# efficient matrix multiplications. Given a Toeplitz matrix with first row
# [t_0, t_1, ... t_{n-1}] and first column [t0, t_{-1}, ..., t_{-(n-1)},
# let C by the circulant matrix with first column [t0, t_{-1}, ...,
# t_{-(n-1)}, 0, t_{n-1}, ..., t_1]. Also adjoin to our input vector `x`
# `n` zeros, to make it a vector of length `2n` (call it y). It can be shown
# that if we take the first n entries of `Cy`, this is equal to the Toeplitz
# multiplication. See:
# http://math.mit.edu/icg/resources/teaching/18.085-spring2015/toeplitz.pdf
# for more details.
x = linalg.adjoint(x) if adjoint_arg else x
expanded_x = array_ops.concat([x, array_ops.zeros_like(x)], axis=-2)
col = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.col)
row = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.row)
circulant_col = array_ops.concat(
[col,
array_ops.zeros_like(col[..., 0:1]),
array_ops.reverse(row[..., 1:], axis=[-1])], axis=-1)
circulant = linear_operator_circulant.LinearOperatorCirculant(
fft_ops.fft(_to_complex(circulant_col)),
input_output_dtype=row.dtype)
result = circulant.matmul(expanded_x, adjoint=adjoint, adjoint_arg=False)
shape = self._shape_tensor(row=row, col=col)
return math_ops.cast(
result[..., :self._domain_dimension_tensor(shape=shape), :],
self.dtype)
def _trace(self):
return math_ops.cast(
self.domain_dimension_tensor(),
dtype=self.dtype) * self.col[..., 0]
def _diag_part(self):
diag_entry = self.col[..., 0:1]
return diag_entry * array_ops.ones(
[self.domain_dimension_tensor()], self.dtype)
def _to_dense(self):
row = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.row)
col = tensor_conversion.convert_to_tensor_v2_with_dispatch(self.col)
total_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(row), array_ops.shape(col))
n = array_ops.shape(row)[-1]
row = array_ops.broadcast_to(row, total_shape)
col = array_ops.broadcast_to(col, total_shape)
# We concatenate the column in reverse order to the row.
# This gives us 2*n + 1 elements.
elements = array_ops.concat(
[array_ops.reverse(col, axis=[-1]), row[..., 1:]], axis=-1)
# Given the above vector, the i-th row of the Toeplitz matrix
# is the last n elements of the above vector shifted i right
# (hence the first row is just the row vector provided, and
# the first element of each row will belong to the column vector).
# We construct these set of indices below.
indices = math_ops.mod(
# How much to shift right. This corresponds to `i`.
math_ops.range(0, n) +
# Specifies the last `n` indices.
math_ops.range(n - 1, -1, -1)[..., array_ops.newaxis],
# Mod out by the total number of elements to ensure the index is
# non-negative (for tf.gather) and < 2 * n - 1.
2 * n - 1)
return array_ops.gather(elements, indices, axis=-1)
@property
def col(self):
return self._col
@property
def row(self):
return self._row
@property
def _composite_tensor_fields(self):
return ("col", "row")
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
return {"col": 1, "row": 1}
def _to_complex(x):
dtype = dtypes.complex64
if x.dtype in [dtypes.float64, dtypes.complex128]:
dtype = dtypes.complex128
return math_ops.cast(x, dtype)
tensorflow/python/ops/linalg/linear_operator_tridiag.py
+401
View File
@@ -0,0 +1,401 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a tridiagonal matrix."""
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import manip_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = ['LinearOperatorTridiag',]
_COMPACT = 'compact'
_MATRIX = 'matrix'
_SEQUENCE = 'sequence'
_DIAGONAL_FORMATS = frozenset({_COMPACT, _MATRIX, _SEQUENCE})
@tf_export('linalg.LinearOperatorTridiag')
@linear_operator.make_composite_tensor
class LinearOperatorTridiag(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] square tridiagonal matrix.
This operator acts like a [batch] square tridiagonal matrix `A` with shape
`[B1,...,Bb, N, N]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x M` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
Example usage:
Create a 3 x 3 tridiagonal linear operator.
>>> superdiag = [3., 4., 5.]
>>> diag = [1., -1., 2.]
>>> subdiag = [6., 7., 8]
>>> operator = tf.linalg.LinearOperatorTridiag(
... [superdiag, diag, subdiag],
... diagonals_format='sequence')
>>> operator.to_dense()
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[ 1., 3., 0.],
[ 7., -1., 4.],
[ 0., 8., 2.]], dtype=float32)>
>>> operator.shape
TensorShape([3, 3])
Scalar Tensor output.
>>> operator.log_abs_determinant()
<tf.Tensor: shape=(), dtype=float32, numpy=4.3307333>
Create a [2, 3] batch of 4 x 4 linear operators.
>>> diagonals = tf.random.normal(shape=[2, 3, 3, 4])
>>> operator = tf.linalg.LinearOperatorTridiag(
... diagonals,
... diagonals_format='compact')
Create a shape [2, 1, 4, 2] vector. Note that this shape is compatible
since the batch dimensions, [2, 1], are broadcast to
operator.batch_shape = [2, 3].
>>> y = tf.random.normal(shape=[2, 1, 4, 2])
>>> x = operator.solve(y)
>>> x
<tf.Tensor: shape=(2, 3, 4, 2), dtype=float32, numpy=...,
dtype=float32)>
#### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, N], with b >= 0
x.shape = [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb].
```
#### Performance
Suppose `operator` is a `LinearOperatorTridiag` of shape `[N, N]`,
and `x.shape = [N, R]`. Then
* `operator.matmul(x)` will take O(N * R) time.
* `operator.solve(x)` will take O(N * R) time.
If instead `operator` and `x` have shape `[B1,...,Bb, N, N]` and
`[B1,...,Bb, N, R]`, every operation increases in complexity by `B1*...*Bb`.
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
diagonals,
diagonals_format=_COMPACT,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name='LinearOperatorTridiag'):
r"""Initialize a `LinearOperatorTridiag`.
Args:
diagonals: `Tensor` or list of `Tensor`s depending on `diagonals_format`.
If `diagonals_format=sequence`, this is a list of three `Tensor`'s each
with shape `[B1, ..., Bb, N]`, `b >= 0, N >= 0`, representing the
superdiagonal, diagonal and subdiagonal in that order. Note the
superdiagonal is padded with an element in the last position, and the
subdiagonal is padded with an element in the front.
If `diagonals_format=matrix` this is a `[B1, ... Bb, N, N]` shaped
`Tensor` representing the full tridiagonal matrix.
If `diagonals_format=compact` this is a `[B1, ... Bb, 3, N]` shaped
`Tensor` with the second to last dimension indexing the
superdiagonal, diagonal and subdiagonal in that order. Note the
superdiagonal is padded with an element in the last position, and the
subdiagonal is padded with an element in the front.
In every case, these `Tensor`s are all floating dtype.
diagonals_format: one of `matrix`, `sequence`, or `compact`. Default is
`compact`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. If `diag.dtype` is real, this is auto-set to `True`.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name for this `LinearOperator`.
Raises:
TypeError: If `diag.dtype` is not an allowed type.
ValueError: If `diag.dtype` is real, and `is_self_adjoint` is not `True`.
"""
parameters = dict(
diagonals=diagonals,
diagonals_format=diagonals_format,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name
)
with ops.name_scope(name, values=[diagonals]):
if diagonals_format not in _DIAGONAL_FORMATS:
raise ValueError(
f'Argument `diagonals_format` must be one of compact, matrix, or '
f'sequence. Received : {diagonals_format}.')
if diagonals_format == _SEQUENCE:
self._diagonals = [linear_operator_util.convert_nonref_to_tensor(
d, name='diag_{}'.format(i)) for i, d in enumerate(diagonals)]
dtype = self._diagonals[0].dtype
else:
self._diagonals = linear_operator_util.convert_nonref_to_tensor(
diagonals, name='diagonals')
dtype = self._diagonals.dtype
self._diagonals_format = diagonals_format
super(LinearOperatorTridiag, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
def _shape(self):
if self.diagonals_format == _MATRIX:
return self.diagonals.shape
if self.diagonals_format == _COMPACT:
# Remove the second to last dimension that contains the value 3.
d_shape = self.diagonals.shape[:-2].concatenate(
self.diagonals.shape[-1])
else:
broadcast_shape = array_ops.broadcast_static_shape(
self.diagonals[0].shape[:-1],
self.diagonals[1].shape[:-1])
broadcast_shape = array_ops.broadcast_static_shape(
broadcast_shape,
self.diagonals[2].shape[:-1])
d_shape = broadcast_shape.concatenate(self.diagonals[1].shape[-1])
return d_shape.concatenate(d_shape[-1])
def _shape_tensor(self, diagonals=None):
diagonals = diagonals if diagonals is not None else self.diagonals
if self.diagonals_format == _MATRIX:
return array_ops.shape(diagonals)
if self.diagonals_format == _COMPACT:
d_shape = array_ops.shape(diagonals[..., 0, :])
else:
broadcast_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(self.diagonals[0])[:-1],
array_ops.shape(self.diagonals[1])[:-1])
broadcast_shape = array_ops.broadcast_dynamic_shape(
broadcast_shape,
array_ops.shape(self.diagonals[2])[:-1])
d_shape = array_ops.concat(
[broadcast_shape, [array_ops.shape(self.diagonals[1])[-1]]], axis=0)
return array_ops.concat([d_shape, [d_shape[-1]]], axis=-1)
def _assert_self_adjoint(self):
# Check the diagonal has non-zero imaginary, and the super and subdiagonals
# are conjugate.
asserts = []
diag_message = (
'This tridiagonal operator contained non-zero '
'imaginary values on the diagonal.')
off_diag_message = (
'This tridiagonal operator has non-conjugate '
'subdiagonal and superdiagonal.')
if self.diagonals_format == _MATRIX:
asserts += [check_ops.assert_equal(
self.diagonals, linalg.adjoint(self.diagonals),
message='Matrix was not equal to its adjoint.')]
elif self.diagonals_format == _COMPACT:
diagonals = tensor_conversion.convert_to_tensor_v2_with_dispatch(
self.diagonals
)
asserts += [linear_operator_util.assert_zero_imag_part(
diagonals[..., 1, :], message=diag_message)]
# Roll the subdiagonal so the shifted argument is at the end.
subdiag = manip_ops.roll(diagonals[..., 2, :], shift=-1, axis=-1)
asserts += [check_ops.assert_equal(
math_ops.conj(subdiag[..., :-1]),
diagonals[..., 0, :-1],
message=off_diag_message)]
else:
asserts += [linear_operator_util.assert_zero_imag_part(
self.diagonals[1], message=diag_message)]
subdiag = manip_ops.roll(self.diagonals[2], shift=-1, axis=-1)
asserts += [check_ops.assert_equal(
math_ops.conj(subdiag[..., :-1]),
self.diagonals[0][..., :-1],
message=off_diag_message)]
return control_flow_ops.group(asserts)
def _construct_adjoint_diagonals(self, diagonals):
# Constructs adjoint tridiagonal matrix from diagonals.
if self.diagonals_format == _SEQUENCE:
diagonals = [math_ops.conj(d) for d in reversed(diagonals)]
# The subdiag and the superdiag swap places, so we need to shift the
# padding argument.
diagonals[0] = manip_ops.roll(diagonals[0], shift=-1, axis=-1)
diagonals[2] = manip_ops.roll(diagonals[2], shift=1, axis=-1)
return diagonals
elif self.diagonals_format == _MATRIX:
return linalg.adjoint(diagonals)
else:
diagonals = math_ops.conj(diagonals)
superdiag, diag, subdiag = array_ops_stack.unstack(
diagonals, num=3, axis=-2)
# The subdiag and the superdiag swap places, so we need
# to shift all arguments.
new_superdiag = manip_ops.roll(subdiag, shift=-1, axis=-1)
new_subdiag = manip_ops.roll(superdiag, shift=1, axis=-1)
return array_ops_stack.stack([new_superdiag, diag, new_subdiag], axis=-2)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
diagonals = self.diagonals
if adjoint:
diagonals = self._construct_adjoint_diagonals(diagonals)
x = linalg.adjoint(x) if adjoint_arg else x
return linalg.tridiagonal_matmul(
diagonals, x,
diagonals_format=self.diagonals_format)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
diagonals = self.diagonals
if adjoint:
diagonals = self._construct_adjoint_diagonals(diagonals)
# TODO(b/144860784): Remove the broadcasting code below once
# tridiagonal_solve broadcasts.
rhs_shape = array_ops.shape(rhs)
k = self._shape_tensor(diagonals)[-1]
broadcast_shape = array_ops.broadcast_dynamic_shape(
self._shape_tensor(diagonals)[:-2], rhs_shape[:-2])
rhs = array_ops.broadcast_to(
rhs, array_ops.concat(
[broadcast_shape, rhs_shape[-2:]], axis=-1))
if self.diagonals_format == _MATRIX:
diagonals = array_ops.broadcast_to(
diagonals, array_ops.concat(
[broadcast_shape, [k, k]], axis=-1))
elif self.diagonals_format == _COMPACT:
diagonals = array_ops.broadcast_to(
diagonals, array_ops.concat(
[broadcast_shape, [3, k]], axis=-1))
else:
diagonals = [
array_ops.broadcast_to(d, array_ops.concat(
[broadcast_shape, [k]], axis=-1)) for d in diagonals]
y = linalg.tridiagonal_solve(
diagonals, rhs,
diagonals_format=self.diagonals_format,
transpose_rhs=adjoint_arg,
conjugate_rhs=adjoint_arg)
return y
def _diag_part(self):
if self.diagonals_format == _MATRIX:
return array_ops.matrix_diag_part(self.diagonals)
elif self.diagonals_format == _SEQUENCE:
diagonal = self.diagonals[1]
return array_ops.broadcast_to(
diagonal, self.shape_tensor()[:-1])
else:
return self.diagonals[..., 1, :]
def _to_dense(self):
if self.diagonals_format == _MATRIX:
return self.diagonals
if self.diagonals_format == _COMPACT:
return gen_array_ops.matrix_diag_v3(
self.diagonals,
k=(-1, 1),
num_rows=-1,
num_cols=-1,
align='LEFT_RIGHT',
padding_value=0.)
diagonals = [
tensor_conversion.convert_to_tensor_v2_with_dispatch(d)
for d in self.diagonals
]
diagonals = array_ops_stack.stack(diagonals, axis=-2)
return gen_array_ops.matrix_diag_v3(
diagonals,
k=(-1, 1),
num_rows=-1,
num_cols=-1,
align='LEFT_RIGHT',
padding_value=0.)
@property
def diagonals(self):
return self._diagonals
@property
def diagonals_format(self):
return self._diagonals_format
@property
def _composite_tensor_fields(self):
return ('diagonals', 'diagonals_format')
@property
def _experimental_parameter_ndims_to_matrix_ndims(self):
diagonal_event_ndims = 2
if self.diagonals_format == _SEQUENCE:
# For the diagonal and the super/sub diagonals.
diagonal_event_ndims = [1, 1, 1]
return {
'diagonals': diagonal_event_ndims,
}
tensorflow/python/ops/linalg/linear_operator_util.py
+627
View File
@@ -0,0 +1,627 @@
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Internal utilities for `LinearOperator` classes."""
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_conversion
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables as variables_module
from tensorflow.python.util import nest
################################################################################
# To make more friendly for TF2.
################################################################################
def convert_nonref_to_tensor(value, dtype=None, dtype_hint=None, name=None):
"""Converts the given `value` to a `Tensor` if input is nonreference type.
This function converts Python objects of various types to `Tensor` objects
except if the input has nonreference semantics. Reference semantics are
characterized by `is_ref` and is any object which is a
`tf.Variable` or instance of `tf.Module`. This function accepts any input
which `tf.convert_to_tensor` would also.
Note: This function diverges from default Numpy behavior for `float` and
`string` types when `None` is present in a Python list or scalar. Rather
than silently converting `None` values, an error will be thrown.
Args:
value: An object whose type has a registered `Tensor` conversion function.
dtype: Optional element type for the returned tensor. If missing, the
type is inferred from the type of `value`.
dtype_hint: Optional element type for the returned tensor,
used when dtype is None. In some cases, a caller may not have a
dtype in mind when converting to a tensor, so dtype_hint
can be used as a soft preference. If the conversion to
`dtype_hint` is not possible, this argument has no effect.
name: Optional name to use if a new `Tensor` is created.
Returns:
tensor: A `Tensor` based on `value`.
Raises:
TypeError: If no conversion function is registered for `value` to `dtype`.
RuntimeError: If a registered conversion function returns an invalid value.
ValueError: If the `value` is a tensor not of given `dtype` in graph mode.
#### Examples:
```python
x = tf.Variable(0.)
y = convert_nonref_to_tensor(x)
x is y
# ==> True
x = tf.constant(0.)
y = convert_nonref_to_tensor(x)
x is y
# ==> True
x = np.array(0.)
y = convert_nonref_to_tensor(x)
x is y
# ==> False
tf.is_tensor(y)
# ==> True
x = tfp.util.DeferredTensor(13.37, lambda x: x)
y = convert_nonref_to_tensor(x)
x is y
# ==> True
tf.is_tensor(y)
# ==> False
tf.equal(y, 13.37)
# ==> True
```
"""
# We explicitly do not use a tf.name_scope to avoid graph clutter.
if value is None:
return None
if is_ref(value):
if dtype is None:
return value
dtype_base = base_dtype(dtype)
value_dtype_base = base_dtype(value.dtype)
if dtype_base != value_dtype_base:
raise TypeError(
f"Argument `value` must be of dtype `{dtype_name(dtype_base)}` "
f"Received: `{dtype_name(value_dtype_base)}`.")
return value
return tensor_conversion.convert_to_tensor_v2_with_dispatch(
value, dtype=dtype, dtype_hint=dtype_hint, name=name
)
def base_dtype(dtype):
"""Returns a non-reference `dtype` based on this `dtype`."""
dtype = dtypes.as_dtype(dtype)
if hasattr(dtype, "base_dtype"):
return dtype.base_dtype
return dtype
def dtype_name(dtype):
"""Returns the string name for this `dtype`."""
dtype = dtypes.as_dtype(dtype)
if hasattr(dtype, "name"):
return dtype.name
if hasattr(dtype, "__name__"):
return dtype.__name__
return str(dtype)
def check_dtype(arg, dtype):
"""Check that arg.dtype == self.dtype."""
if arg.dtype.base_dtype != dtype:
raise TypeError(
f"Expected argument to have dtype {dtype}. Found: {arg.dtype} in "
f"tensor {arg}.")
def is_ref(x):
"""Evaluates if the object has reference semantics.
An object is deemed "reference" if it is a `tf.Variable` instance or is
derived from a `tf.Module` with `dtype` and `shape` properties.
Args:
x: Any object.
Returns:
is_ref: Python `bool` indicating input is has nonreference semantics, i.e.,
is a `tf.Variable` or a `tf.Module` with `dtype` and `shape` properties.
"""
return (
# Note: we check that tf.Variable is a class because we might be using a
# different backend other than TF.
isinstance(x, variables_module.Variable) or
(isinstance(x, module.Module) and hasattr(x, "dtype") and
hasattr(x, "shape")))
def assert_not_ref_type(x, arg_name):
if is_ref(x):
raise TypeError(
f"Argument {arg_name} cannot be reference type. Found: {type(x)}.")
################################################################################
# Asserts.
################################################################################
def assert_no_entries_with_modulus_zero(
x, message=None, name="assert_no_entries_with_modulus_zero"):
"""Returns `Op` that asserts Tensor `x` has no entries with modulus zero.
Args:
x: Numeric `Tensor`, real, integer, or complex.
message: A string message to prepend to failure message.
name: A name to give this `Op`.
Returns:
An `Op` that asserts `x` has no entries with modulus zero.
"""
with ops.name_scope(name, values=[x]):
x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name="x")
dtype = x.dtype.base_dtype
should_be_nonzero = math_ops.abs(x)
zero = tensor_conversion.convert_to_tensor_v2_with_dispatch(
0, dtype=dtype.real_dtype
)
return check_ops.assert_less(zero, should_be_nonzero, message=message)
def assert_zero_imag_part(x, message=None, name="assert_zero_imag_part"):
"""Returns `Op` that asserts Tensor `x` has no non-zero imaginary parts.
Args:
x: Numeric `Tensor`, real, integer, or complex.
message: A string message to prepend to failure message.
name: A name to give this `Op`.
Returns:
An `Op` that asserts `x` has no entries with modulus zero.
"""
with ops.name_scope(name, values=[x]):
x = tensor_conversion.convert_to_tensor_v2_with_dispatch(x, name="x")
dtype = x.dtype.base_dtype
if dtype.is_floating:
return control_flow_ops.no_op()
zero = tensor_conversion.convert_to_tensor_v2_with_dispatch(
0, dtype=dtype.real_dtype
)
return check_ops.assert_equal(zero, math_ops.imag(x), message=message)
def assert_compatible_matrix_dimensions(operator, x):
"""Assert that an argument to solve/matmul has proper domain dimension.
If `operator.shape[-2:] = [M, N]`, and `x.shape[-2:] = [Q, R]`, then
`operator.matmul(x)` is defined only if `N = Q`. This `Op` returns an
`Assert` that "fires" if this is not the case. Static checks are already
done by the base class `LinearOperator`.
Args:
operator: `LinearOperator`.
x: `Tensor`.
Returns:
`Assert` `Op`.
"""
# Static checks are done in the base class. Only tensor asserts here.
assert_same_dd = check_ops.assert_equal(
array_ops.shape(x)[-2],
operator.domain_dimension_tensor(),
# This error message made to look similar to error raised by static check
# in the base class.
message=("Dimensions are not compatible. "
"shape[-2] of argument to be the same as this operator"))
return assert_same_dd
def assert_is_batch_matrix(tensor):
"""Static assert that `tensor` has rank `2` or higher."""
sh = tensor.shape
if sh.ndims is not None and sh.ndims < 2:
raise ValueError(
f"Expected [batch] matrix to have at least two dimensions. Found: "
f"{tensor}.")
def shape_tensor(shape, name=None):
"""Convert Tensor using default type, unless empty list or tuple."""
# Works just like random_ops._ShapeTensor.
if isinstance(shape, (tuple, list)) and not shape:
dtype = dtypes.int32
else:
dtype = None
return tensor_conversion.convert_to_tensor_v2_with_dispatch(
shape, dtype=dtype, name=name
)
################################################################################
# Broadcasting versions of common linear algebra functions.
# TODO(b/77519145) Do this more efficiently in some special cases.
################################################################################
def broadcast_matrix_batch_dims(batch_matrices, name=None):
"""Broadcast leading dimensions of zero or more [batch] matrices.
Example broadcasting one batch dim of two simple matrices.
```python
x = [[1, 2],
[3, 4]] # Shape [2, 2], no batch dims
y = [[[1]]] # Shape [1, 1, 1], 1 batch dim of shape [1]
x_bc, y_bc = broadcast_matrix_batch_dims([x, y])
x_bc
==> [[[1, 2],
[3, 4]]] # Shape [1, 2, 2], 1 batch dim of shape [1].
y_bc
==> same as y
```
Example broadcasting many batch dims
```python
x = tf.random.normal(shape=(2, 3, 1, 4, 4))
y = tf.random.normal(shape=(1, 3, 2, 5, 5))
x_bc, y_bc = broadcast_matrix_batch_dims([x, y])
x_bc.shape
==> (2, 3, 2, 4, 4)
y_bc.shape
==> (2, 3, 2, 5, 5)
```
Args:
batch_matrices: Iterable of `Tensor`s, each having two or more dimensions.
name: A string name to prepend to created ops.
Returns:
bcast_matrices: List of `Tensor`s, with `bcast_matrices[i]` containing
the values from `batch_matrices[i]`, with possibly broadcast batch dims.
Raises:
ValueError: If any input `Tensor` is statically determined to have less
than two dimensions.
"""
with ops.name_scope(
name or "broadcast_matrix_batch_dims", values=batch_matrices):
check_ops.assert_proper_iterable(batch_matrices)
batch_matrices = list(batch_matrices)
for i, mat in enumerate(batch_matrices):
batch_matrices[i] = tensor_conversion.convert_to_tensor_v2_with_dispatch(
mat
)
assert_is_batch_matrix(batch_matrices[i])
if len(batch_matrices) < 2:
return batch_matrices
# Try static broadcasting.
# bcast_batch_shape is the broadcast batch shape of ALL matrices.
# E.g. if batch_matrices = [x, y], with
# x.shape = [2, j, k] (batch shape = [2])
# y.shape = [3, 1, l, m] (batch shape = [3, 1])
# ==> bcast_batch_shape = [3, 2]
bcast_batch_shape = batch_matrices[0].shape[:-2]
for mat in batch_matrices[1:]:
bcast_batch_shape = array_ops.broadcast_static_shape(
bcast_batch_shape,
mat.shape[:-2])
if bcast_batch_shape.is_fully_defined():
for i, mat in enumerate(batch_matrices):
if mat.shape[:-2] != bcast_batch_shape:
bcast_shape = array_ops.concat(
[bcast_batch_shape.as_list(), array_ops.shape(mat)[-2:]], axis=0)
batch_matrices[i] = array_ops.broadcast_to(mat, bcast_shape)
return batch_matrices
# Since static didn't work, do dynamic, which always copies data.
bcast_batch_shape = array_ops.shape(batch_matrices[0])[:-2]
for mat in batch_matrices[1:]:
bcast_batch_shape = array_ops.broadcast_dynamic_shape(
bcast_batch_shape,
array_ops.shape(mat)[:-2])
for i, mat in enumerate(batch_matrices):
batch_matrices[i] = array_ops.broadcast_to(
mat,
array_ops.concat(
[bcast_batch_shape, array_ops.shape(mat)[-2:]], axis=0))
return batch_matrices
def matrix_solve_with_broadcast(matrix, rhs, adjoint=False, name=None):
"""Solve systems of linear equations."""
with ops.name_scope(name, "MatrixSolveWithBroadcast", [matrix, rhs]):
matrix = tensor_conversion.convert_to_tensor_v2_with_dispatch(
matrix, name="matrix"
)
rhs = tensor_conversion.convert_to_tensor_v2_with_dispatch(
rhs, name="rhs", dtype=matrix.dtype
)
# If either matrix/rhs has extra dims, we can reshape to get rid of them.
matrix, rhs, reshape_inv, still_need_to_transpose = _reshape_for_efficiency(
matrix, rhs, adjoint_a=adjoint)
# This will broadcast by brute force if we still need to.
matrix, rhs = broadcast_matrix_batch_dims([matrix, rhs])
solution = linalg_ops.matrix_solve(
matrix, rhs, adjoint=adjoint and still_need_to_transpose)
return reshape_inv(solution)
def _reshape_for_efficiency(a,
b,
transpose_a=False,
transpose_b=False,
adjoint_a=False,
adjoint_b=False):
"""Maybe reshape a, b, and return an inverse map. For matmul/solve."""
def identity(x):
return x
# At this point, we have not taken transpose/adjoint of a/b.
still_need_to_transpose = True
if a.shape.ndims is None or b.shape.ndims is None:
return a, b, identity, still_need_to_transpose
# This could be handled in the future, but seems less common.
if a.shape.ndims >= b.shape.ndims:
return a, b, identity, still_need_to_transpose
# From now on, we might modify b, but will not modify a.
# Suppose:
# a.shape = C + [m, n], b.shape =
# b.shape = S + C + [n, r]
b_extra_ndims = b.shape.ndims - a.shape.ndims
# b_extra_sh = S, b_main_sh = C + [n, r]
b_extra_sh = array_ops.shape(b)[:b_extra_ndims]
b_main_sh = array_ops.shape(b)[b_extra_ndims:]
# No reason to flip unless the extra dims of b are big enough. Why?
# Assume adjoint/transpose = False. Then...
# By not flipping, we have to replicate a to shape
# b_extra_sh + a.shape,
# which could use extra memory. But in all cases, the final output has shape
# b_extra_sh + a.shape[:-1] + [b.shape[-1]]
# So we only end up creating a larger object if the end dim of b is smaller
# than the end dim of a. This often happens, e.g. if b was a vector that was
# expanded to a matrix (by appending a singleton).
# Since adjoint/transpose may not be False, we must make adjustments here.
# The dim of b that holds the multiple equations.
a_domain_sz_ = a.shape[-2 if adjoint_a or transpose_a else -1]
b_eq_sz_ = b.shape[-2 if adjoint_b or transpose_b else -1]
b_extra_sz_ = (
np.prod(b.shape[:b_extra_ndims].as_list())
if b.shape[:b_extra_ndims].is_fully_defined() else None)
if (a_domain_sz_ is not None and b_eq_sz_ is not None and
b_extra_sz_ is not None):
if b_extra_sz_ < 2 or a_domain_sz_ <= b_eq_sz_:
return a, b, identity, still_need_to_transpose
# At this point, we're flipping for sure!
# Any transposes/adjoints will happen here explicitly, rather than in calling
# code. Why? To avoid having to write separate complex code for each case.
if adjoint_a:
a = array_ops.matrix_transpose(a, conjugate=True)
elif transpose_a:
a = array_ops.matrix_transpose(a, conjugate=False)
if adjoint_b:
b = array_ops.matrix_transpose(b, conjugate=True)
elif transpose_a:
b = array_ops.matrix_transpose(b, conjugate=False)
still_need_to_transpose = False
# Recompute shapes, since the transpose/adjoint may have changed them.
b_extra_sh = array_ops.shape(b)[:b_extra_ndims]
b_main_sh = array_ops.shape(b)[b_extra_ndims:]
# Permutation to put the extra dims at the end.
perm = (
np.concatenate(
(np.arange(b_extra_ndims, b.shape.ndims),
np.arange(0, b_extra_ndims)), 0))
b_extra_on_end = array_ops.transpose(b, perm=perm)
# Now squash this end into one long dim.
b_squashed_end = array_ops.reshape(
b_extra_on_end, array_ops.concat((b_main_sh[:-1], [-1]), 0))
def reshape_inv(y):
# Expand the extra dims hanging off the end, "b_extra_sh".
# Note we use y_sh[:-1] + [b_main_sh[-1]] rather than b_main_sh, because y
# Could have different batch dims than a and b, because of broadcasting.
y_extra_shape = array_ops.concat(
(array_ops.shape(y)[:-1], [b_main_sh[-1]], b_extra_sh), 0)
y_extra_on_end = array_ops.reshape(y, y_extra_shape)
inverse_perm = np.argsort(perm)
return array_ops.transpose(y_extra_on_end, perm=inverse_perm)
return a, b_squashed_end, reshape_inv, still_need_to_transpose
################################################################################
# Helpers for hints.
################################################################################
def is_adjoint_pair(x, y):
"""True iff x and y are adjoints of each other (by id, not entries)."""
if x is y: # Note that if x is y then all of their hints are the same!
if x.is_self_adjoint is False: # pylint:disable=g-bool-id-comparison
return False
if x.is_self_adjoint:
return True
# Use the fact that if x = LinearOperatorAdjoint(y), then x.H is y.
return x.H is y or y.H is x
def is_aat_form(operators):
"""Returns True if operators is of the form A @ A.H, possibly recursively."""
operators = list(operators)
if not operators:
raise ValueError("AAT form is undefined for empty operators")
if len(operators) % 2:
return False
# Check for forms like (A1 @ A2) @ (A2.H @ A1.H)
return all(
is_adjoint_pair(operators[i], operators[-1 - i])
for i in range(len(operators) // 2))
def use_operator_or_provided_hint_unless_contradicting(
operator, hint_attr_name, provided_hint_value, message):
"""Get combined hint in the case where operator.hint should equal hint.
Args:
operator: LinearOperator that a meta-operator was initialized with.
hint_attr_name: String name for the attribute.
provided_hint_value: Bool or None. Value passed by user in initialization.
message: Error message to print if hints contradict.
Returns:
True, False, or None.
Raises:
ValueError: If hints contradict.
"""
op_hint = getattr(operator, hint_attr_name)
# pylint: disable=g-bool-id-comparison
if op_hint is False and provided_hint_value:
raise ValueError(message)
if op_hint and provided_hint_value is False:
raise ValueError(message)
if op_hint or provided_hint_value:
return True
if op_hint is False or provided_hint_value is False:
return False
# pylint: enable=g-bool-id-comparison
return None
################################################################################
# Utilities for blockwise operators.
################################################################################
def arg_is_blockwise(block_dimensions, arg, arg_split_dim):
"""Detect if input should be interpreted as a list of blocks."""
# Tuples and lists of length equal to the number of operators may be
# blockwise.
if (isinstance(arg, (tuple, list)) and len(arg) == len(block_dimensions)):
# If the elements of the iterable are not nested, interpret the input as
# blockwise.
if not any(nest.is_nested(x) for x in arg):
return True
else:
arg_dims = [
tensor_conversion.convert_to_tensor_v2_with_dispatch(x).shape[
arg_split_dim
]
for x in arg
]
self_dims = [dim.value for dim in block_dimensions]
# If none of the operator dimensions are known, interpret the input as
# blockwise if its matching dimensions are unequal.
if all(self_d is None for self_d in self_dims):
# A nested tuple/list with a single outermost element is not blockwise
if len(arg_dims) == 1:
return False
elif any(dim != arg_dims[0] for dim in arg_dims):
return True
else:
raise ValueError(
"Parsing of the input structure is ambiguous. Please input "
"a blockwise iterable of `Tensor`s or a single `Tensor`.")
# If input dimensions equal the respective (known) blockwise operator
# dimensions, then the input is blockwise.
if all(self_d == arg_d or self_d is None
for self_d, arg_d in zip(self_dims, arg_dims)):
return True
# If input dimensions equals are all equal, and are greater than or equal
# to the sum of the known operator dimensions, interpret the input as
# blockwise.
# input is not blockwise.
self_dim = sum(self_d for self_d in self_dims if self_d is not None)
if all(s == arg_dims[0] for s in arg_dims) and arg_dims[0] >= self_dim:
return False
# If none of these conditions is met, the input shape is mismatched.
raise ValueError("Input dimension does not match operator dimension.")
else:
return False
def split_arg_into_blocks(block_dims, block_dims_fn, arg, axis=-1):
"""Split `x` into blocks matching `operators`'s `domain_dimension`.
Specifically, if we have a blockwise lower-triangular matrix, with block
sizes along the diagonal `[M_j, M_j] j = 0,1,2..J`, this method splits `arg`
on `axis` into `J` tensors, whose shape at `axis` is `M_j`.
Args:
block_dims: Iterable of `TensorShapes`.
block_dims_fn: Callable returning an iterable of `Tensor`s.
arg: `Tensor`. `arg` is split into `J` tensors.
axis: Python `Integer` representing the axis to split `arg` on.
Returns:
A list of `Tensor`s.
"""
block_sizes = [dim.value for dim in block_dims]
if any(d is None for d in block_sizes):
block_sizes = block_dims_fn()
return array_ops.split(arg, block_sizes, axis=axis)
tensorflow/python/ops/linalg/linear_operator_zeros.py
+500
View File
@@ -0,0 +1,500 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`LinearOperator` acting like a zero matrix."""
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import check_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.ops.linalg import linear_operator
from tensorflow.python.ops.linalg import linear_operator_util
from tensorflow.python.util.tf_export import tf_export
__all__ = [
"LinearOperatorZeros",
]
@tf_export("linalg.LinearOperatorZeros")
@linear_operator.make_composite_tensor
class LinearOperatorZeros(linear_operator.LinearOperator):
"""`LinearOperator` acting like a [batch] zero matrix.
This operator acts like a [batch] zero matrix `A` with shape
`[B1,...,Bb, N, M]` for some `b >= 0`. The first `b` indices index a
batch member. For every batch index `(i1,...,ib)`, `A[i1,...,ib, : :]` is
an `N x M` matrix. This matrix `A` is not materialized, but for
purposes of broadcasting this shape will be relevant.
`LinearOperatorZeros` is initialized with `num_rows`, and optionally
`num_columns, `batch_shape`, and `dtype` arguments. If `num_columns` is
`None`, then this operator will be initialized as a square matrix. If
`batch_shape` is `None`, this operator efficiently passes through all
arguments. If `batch_shape` is provided, broadcasting may occur, which will
require making copies.
```python
# Create a 2 x 2 zero matrix.
operator = LinearOperatorZero(num_rows=2, dtype=tf.float32)
operator.to_dense()
==> [[0., 0.]
[0., 0.]]
operator.shape
==> [2, 2]
operator.determinant()
==> 0.
x = ... Shape [2, 4] Tensor
operator.matmul(x)
==> Shape [2, 4] Tensor, same as x.
# Create a 2-batch of 2x2 zero matrices
operator = LinearOperatorZeros(num_rows=2, batch_shape=[2])
operator.to_dense()
==> [[[0., 0.]
[0., 0.]],
[[0., 0.]
[0., 0.]]]
# Here, even though the operator has a batch shape, the input is the same as
# the output, so x can be passed through without a copy. The operator is able
# to detect that no broadcast is necessary because both x and the operator
# have statically defined shape.
x = ... Shape [2, 2, 3]
operator.matmul(x)
==> Shape [2, 2, 3] Tensor, same as tf.zeros_like(x)
# Here the operator and x have different batch_shape, and are broadcast.
# This requires a copy, since the output is different size than the input.
x = ... Shape [1, 2, 3]
operator.matmul(x)
==> Shape [2, 2, 3] Tensor, equal to tf.zeros_like([x, x])
```
### Shape compatibility
This operator acts on [batch] matrix with compatible shape.
`x` is a batch matrix with compatible shape for `matmul` and `solve` if
```
operator.shape = [B1,...,Bb] + [N, M], with b >= 0
x.shape = [C1,...,Cc] + [M, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb] to [D1,...,Dd]
```
#### Matrix property hints
This `LinearOperator` is initialized with boolean flags of the form `is_X`,
for `X = non_singular, self_adjoint, positive_definite, square`.
These have the following meaning:
* If `is_X == True`, callers should expect the operator to have the
property `X`. This is a promise that should be fulfilled, but is *not* a
runtime assert. For example, finite floating point precision may result
in these promises being violated.
* If `is_X == False`, callers should expect the operator to not have `X`.
* If `is_X == None` (the default), callers should have no expectation either
way.
"""
def __init__(self,
num_rows,
num_columns=None,
batch_shape=None,
dtype=None,
is_non_singular=False,
is_self_adjoint=True,
is_positive_definite=False,
is_square=True,
assert_proper_shapes=False,
name="LinearOperatorZeros"):
r"""Initialize a `LinearOperatorZeros`.
The `LinearOperatorZeros` is initialized with arguments defining `dtype`
and shape.
This operator is able to broadcast the leading (batch) dimensions, which
sometimes requires copying data. If `batch_shape` is `None`, the operator
can take arguments of any batch shape without copying. See examples.
Args:
num_rows: Scalar non-negative integer `Tensor`. Number of rows in the
corresponding zero matrix.
num_columns: Scalar non-negative integer `Tensor`. Number of columns in
the corresponding zero matrix. If `None`, defaults to the value of
`num_rows`.
batch_shape: Optional `1-D` integer `Tensor`. The shape of the leading
dimensions. If `None`, this operator has no leading dimensions.
dtype: Data type of the matrix that this operator represents.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
assert_proper_shapes: Python `bool`. If `False`, only perform static
checks that initialization and method arguments have proper shape.
If `True`, and static checks are inconclusive, add asserts to the graph.
name: A name for this `LinearOperator`
Raises:
ValueError: If `num_rows` is determined statically to be non-scalar, or
negative.
ValueError: If `num_columns` is determined statically to be non-scalar,
or negative.
ValueError: If `batch_shape` is determined statically to not be 1-D, or
negative.
ValueError: If any of the following is not `True`:
`{is_self_adjoint, is_non_singular, is_positive_definite}`.
"""
parameters = dict(
num_rows=num_rows,
num_columns=num_columns,
batch_shape=batch_shape,
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
assert_proper_shapes=assert_proper_shapes,
name=name
)
dtype = dtype or dtypes.float32
self._assert_proper_shapes = assert_proper_shapes
with ops.name_scope(name):
dtype = dtypes.as_dtype(dtype)
if not is_self_adjoint and is_square:
raise ValueError("A zero operator is always self adjoint.")
if is_non_singular:
raise ValueError("A zero operator is always singular.")
if is_positive_definite:
raise ValueError("A zero operator is always not positive-definite.")
super(LinearOperatorZeros, self).__init__(
dtype=dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
parameters=parameters,
name=name)
linear_operator_util.assert_not_ref_type(num_rows, "num_rows")
linear_operator_util.assert_not_ref_type(num_columns, "num_columns")
linear_operator_util.assert_not_ref_type(batch_shape, "batch_shape")
self._num_rows = linear_operator_util.shape_tensor(
num_rows, name="num_rows")
self._num_rows_static = tensor_util.constant_value(self._num_rows)
if num_columns is None:
num_columns = num_rows
self._num_columns = linear_operator_util.shape_tensor(
num_columns, name="num_columns")
self._num_columns_static = tensor_util.constant_value(self._num_columns)
self._check_domain_range_possibly_add_asserts()
if (self._num_rows_static is not None and
self._num_columns_static is not None):
if is_square and self._num_rows_static != self._num_columns_static:
raise ValueError(
"LinearOperatorZeros initialized as is_square=True, but got "
"num_rows({}) != num_columns({})".format(
self._num_rows_static,
self._num_columns_static))
if batch_shape is None:
self._batch_shape_arg = None
else:
self._batch_shape_arg = linear_operator_util.shape_tensor(
batch_shape, name="batch_shape_arg")
self._batch_shape_static = tensor_util.constant_value(
self._batch_shape_arg)
self._check_batch_shape_possibly_add_asserts()
def _shape(self):
matrix_shape = tensor_shape.TensorShape((self._num_rows_static,
self._num_columns_static))
if self._batch_shape_arg is None:
return matrix_shape
batch_shape = tensor_shape.TensorShape(self._batch_shape_static)
return batch_shape.concatenate(matrix_shape)
def _shape_tensor(self):
matrix_shape = array_ops_stack.stack(
(self._num_rows, self._num_columns), axis=0)
if self._batch_shape_arg is None:
return matrix_shape
return array_ops.concat((self._batch_shape_arg, matrix_shape), 0)
def _assert_non_singular(self):
raise errors.InvalidArgumentError(
node_def=None, op=None, message="Zero operators are always "
"non-invertible.")
def _assert_positive_definite(self):
raise errors.InvalidArgumentError(
node_def=None, op=None, message="Zero operators are always "
"non-positive definite.")
def _assert_self_adjoint(self):
return control_flow_ops.no_op("assert_self_adjoint")
def _possibly_broadcast_batch_shape(self, x):
"""Return 'x', possibly after broadcasting the leading dimensions."""
# If we have no batch shape, our batch shape broadcasts with everything!
if self._batch_shape_arg is None:
return x
# Static attempt:
# If we determine that no broadcast is necessary, pass x through
# If we need a broadcast, add to an array of zeros.
#
# special_shape is the shape that, when broadcast with x's shape, will give
# the correct broadcast_shape. Note that
# We have already verified the second to last dimension of self.shape
# matches x's shape in assert_compatible_matrix_dimensions.
# Also, the final dimension of 'x' can have any shape.
# Therefore, the final two dimensions of special_shape are 1's.
special_shape = self.batch_shape.concatenate([1, 1])
bshape = array_ops.broadcast_static_shape(x.shape, special_shape)
if special_shape.is_fully_defined():
# bshape.is_fully_defined iff special_shape.is_fully_defined.
if bshape == x.shape:
return x
# Use the built in broadcasting of addition.
zeros = array_ops.zeros(shape=special_shape, dtype=self.dtype)
return x + zeros
# Dynamic broadcast:
# Always add to an array of zeros, rather than using a "cond", since a
# cond would require copying data from GPU --> CPU.
special_shape = array_ops.concat((self.batch_shape_tensor(), [1, 1]), 0)
zeros = array_ops.zeros(shape=special_shape, dtype=self.dtype)
return x + zeros
def _matmul(self, x, adjoint=False, adjoint_arg=False):
if self._assert_proper_shapes:
x = linalg.adjoint(x) if adjoint_arg else x
aps = linear_operator_util.assert_compatible_matrix_dimensions(self, x)
x = control_flow_ops.with_dependencies([aps], x)
if self.is_square:
# Note that adjoint has no effect since this matrix is self-adjoint.
if adjoint_arg:
output_shape = array_ops.concat([
array_ops.shape(x)[:-2],
[array_ops.shape(x)[-1], array_ops.shape(x)[-2]]], axis=0)
else:
output_shape = array_ops.shape(x)
return self._possibly_broadcast_batch_shape(
array_ops.zeros(shape=output_shape, dtype=x.dtype))
x_shape = array_ops.shape(x)
n = self._num_columns if adjoint else self._num_rows
m = x_shape[-2] if adjoint_arg else x_shape[-1]
output_shape = array_ops.concat([x_shape[:-2], [n, m]], axis=0)
zeros = array_ops.zeros(shape=output_shape, dtype=x.dtype)
return self._possibly_broadcast_batch_shape(zeros)
def _linop_matmul(
self,
left_operator: "LinearOperatorZeros",
right_operator: linear_operator.LinearOperator
) -> linear_operator.LinearOperator:
if not left_operator.is_square or not right_operator.is_square:
raise ValueError("Matmul with non-square `LinearOperator`s or non-square "
"`LinearOperatorZeros` not supported at this time.")
return left_operator
def _determinant(self):
if self.batch_shape.is_fully_defined():
return array_ops.zeros(shape=self.batch_shape, dtype=self.dtype)
else:
return array_ops.zeros(shape=self.batch_shape_tensor(), dtype=self.dtype)
def _trace(self):
# Get Tensor of all zeros of same shape as self.batch_shape.
if self.batch_shape.is_fully_defined():
return array_ops.zeros(shape=self.batch_shape, dtype=self.dtype)
else:
return array_ops.zeros(shape=self.batch_shape_tensor(), dtype=self.dtype)
def _diag_part(self):
return self._zeros_diag()
def add_to_tensor(self, mat, name="add_to_tensor"):
"""Add matrix represented by this operator to `mat`. Equiv to `I + mat`.
Args:
mat: `Tensor` with same `dtype` and shape broadcastable to `self`.
name: A name to give this `Op`.
Returns:
A `Tensor` with broadcast shape and same `dtype` as `self`.
"""
return self._possibly_broadcast_batch_shape(mat)
def _check_domain_range_possibly_add_asserts(self):
"""Static check of init arg `num_rows`, possibly add asserts."""
# Possibly add asserts.
if self._assert_proper_shapes:
self._num_rows = control_flow_ops.with_dependencies([
check_ops.assert_rank(
self._num_rows,
0,
message="Argument num_rows must be a 0-D Tensor."),
check_ops.assert_non_negative(
self._num_rows,
message="Argument num_rows must be non-negative."),
], self._num_rows)
self._num_columns = control_flow_ops.with_dependencies([
check_ops.assert_rank(
self._num_columns,
0,
message="Argument num_columns must be a 0-D Tensor."),
check_ops.assert_non_negative(
self._num_columns,
message="Argument num_columns must be non-negative."),
], self._num_columns)
# Static checks.
if not self._num_rows.dtype.is_integer:
raise TypeError("Argument num_rows must be integer type. Found:"
" %s" % self._num_rows)
if not self._num_columns.dtype.is_integer:
raise TypeError("Argument num_columns must be integer type. Found:"
" %s" % self._num_columns)
num_rows_static = self._num_rows_static
num_columns_static = self._num_columns_static
if num_rows_static is not None:
if num_rows_static.ndim != 0:
raise ValueError("Argument num_rows must be a 0-D Tensor. Found:"
" %s" % num_rows_static)
if num_rows_static < 0:
raise ValueError("Argument num_rows must be non-negative. Found:"
" %s" % num_rows_static)
if num_columns_static is not None:
if num_columns_static.ndim != 0:
raise ValueError("Argument num_columns must be a 0-D Tensor. Found:"
" %s" % num_columns_static)
if num_columns_static < 0:
raise ValueError("Argument num_columns must be non-negative. Found:"
" %s" % num_columns_static)
def _check_batch_shape_possibly_add_asserts(self):
"""Static check of init arg `batch_shape`, possibly add asserts."""
if self._batch_shape_arg is None:
return
# Possibly add asserts
if self._assert_proper_shapes:
self._batch_shape_arg = control_flow_ops.with_dependencies([
check_ops.assert_rank(
self._batch_shape_arg,
1,
message="Argument batch_shape must be a 1-D Tensor."),
check_ops.assert_non_negative(
self._batch_shape_arg,
message="Argument batch_shape must be non-negative."),
], self._batch_shape_arg)
# Static checks
if not self._batch_shape_arg.dtype.is_integer:
raise TypeError("Argument batch_shape must be integer type. Found:"
" %s" % self._batch_shape_arg)
if self._batch_shape_static is None:
return # Cannot do any other static checks.
if self._batch_shape_static.ndim != 1:
raise ValueError("Argument batch_shape must be a 1-D Tensor. Found:"
" %s" % self._batch_shape_static)
if np.any(self._batch_shape_static < 0):
raise ValueError("Argument batch_shape must be non-negative. Found:"
"%s" % self._batch_shape_static)
def _min_matrix_dim(self):
"""Minimum of domain/range dimension, if statically available, else None."""
domain_dim = self.domain_dimension.value
range_dim = self.range_dimension.value
if domain_dim is None or range_dim is None:
return None
return min(domain_dim, range_dim)
def _min_matrix_dim_tensor(self):
"""Minimum of domain/range dimension, as a tensor."""
return math_ops.reduce_min(self.shape_tensor()[-2:])
def _zeros_diag(self):
"""Returns the diagonal of this operator as all zeros."""
if self.shape.is_fully_defined():
d_shape = self.batch_shape.concatenate([self._min_matrix_dim()])
else:
d_shape = array_ops.concat(
[self.batch_shape_tensor(),
[self._min_matrix_dim_tensor()]], axis=0)
return array_ops.zeros(shape=d_shape, dtype=self.dtype)
def _eigvals(self):
return self._zeros_diag()
@property
def _composite_tensor_prefer_static_fields(self):
return ("num_rows", "num_columns", "batch_shape")
@property
def _composite_tensor_fields(self):
return ("num_rows", "num_columns", "batch_shape", "dtype",
"assert_proper_shapes")
def __getitem__(self, slices):
# Slice the batch shape and return a new LinearOperatorIdentity.
# Use a proxy shape and slice it. Use this as the new batch shape
new_batch_shape = array_ops.shape(
array_ops.ones(self._batch_shape_arg)[slices])
parameters = dict(self.parameters, batch_shape=new_batch_shape)
return LinearOperatorZeros(**parameters)
tensorflow/python/ops/linalg/property_hint_util.py
+87
View File
@@ -0,0 +1,87 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Common utilities for LinearOperator property hints."""
# Note: only use this method in the commuting case.
def combined_commuting_self_adjoint_hint(operator_a, operator_b):
"""Get combined hint for self-adjoint-ness."""
# The property is preserved under composition when the operators commute.
if operator_a.is_self_adjoint and operator_b.is_self_adjoint:
return True
# The property is not preserved when an operator with the property is composed
# with an operator without the property.
# pylint:disable=g-bool-id-comparison
if ((operator_a.is_self_adjoint is True and
operator_b.is_self_adjoint is False) or
(operator_a.is_self_adjoint is False and
operator_b.is_self_adjoint is True)):
return False
# pylint:enable=g-bool-id-comparison
# The property is not known when operators are not known to have the property
# or both operators don't have the property (the property for the complement
# class is not closed under composition).
return None
def is_square(operator_a, operator_b):
"""Return a hint to whether the composition is square."""
if operator_a.is_square and operator_b.is_square:
return True
if operator_a.is_square is False and operator_b.is_square is False: # pylint:disable=g-bool-id-comparison
# Let A have shape [B, M, N], B have shape [B, N, L].
m = operator_a.range_dimension
l = operator_b.domain_dimension
if m is not None and l is not None:
return m == l
if (operator_a.is_square != operator_b.is_square) and (
operator_a.is_square is not None and operator_b.is_square is not None):
return False
return None
# Note: Positive definiteness is only guaranteed to be preserved
# when the operators commute and are symmetric. Only use this method in
# commuting cases.
def combined_commuting_positive_definite_hint(operator_a, operator_b):
"""Get combined PD hint for compositions."""
# pylint:disable=g-bool-id-comparison
if (operator_a.is_positive_definite is True and
operator_a.is_self_adjoint is True and
operator_b.is_positive_definite is True and
operator_b.is_self_adjoint is True):
return True
# pylint:enable=g-bool-id-comparison
return None
def combined_non_singular_hint(operator_a, operator_b):
"""Get combined hint for when ."""
# If either operator is not-invertible the composition isn't.
# pylint:disable=g-bool-id-comparison
if (operator_a.is_non_singular is False or
operator_b.is_non_singular is False):
return False
# pylint:enable=g-bool-id-comparison
return operator_a.is_non_singular and operator_b.is_non_singular
tensorflow/python/ops/linalg/slicing.py
+184
View File
@@ -0,0 +1,184 @@
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for slicing in to a `LinearOperator`."""
import collections
import functools
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.util import nest
__all__ = ['batch_slice']
def _prefer_static_where(condition, x, y):
args = [condition, x, y]
constant_args = [tensor_util.constant_value(a) for a in args]
# Do this statically.
if all(arg is not None for arg in constant_args):
condition_, x_, y_ = constant_args
return np.where(condition_, x_, y_)
return array_ops.where(condition, x, y)
def _broadcast_parameter_with_batch_shape(
param, param_ndims_to_matrix_ndims, batch_shape):
"""Broadcasts `param` with the given batch shape, recursively."""
if hasattr(param, 'batch_shape_tensor'):
# Recursively broadcast every parameter inside the operator.
override_dict = {}
for name, ndims in param._experimental_parameter_ndims_to_matrix_ndims.items(): # pylint:disable=protected-access,line-too-long
sub_param = getattr(param, name)
override_dict[name] = nest.map_structure_up_to(
sub_param, functools.partial(
_broadcast_parameter_with_batch_shape,
batch_shape=batch_shape), sub_param, ndims)
parameters = dict(param.parameters, **override_dict)
return type(param)(**parameters)
base_shape = array_ops.concat(
[batch_shape, array_ops.ones(
[param_ndims_to_matrix_ndims], dtype=dtypes.int32)], axis=0)
return array_ops.broadcast_to(
param,
array_ops.broadcast_dynamic_shape(base_shape, array_ops.shape(param)))
def _sanitize_slices(slices, intended_shape, deficient_shape):
"""Restricts slices to avoid overflowing size-1 (broadcast) dimensions.
Args:
slices: iterable of slices received by `__getitem__`.
intended_shape: int `Tensor` shape for which the slices were intended.
deficient_shape: int `Tensor` shape to which the slices will be applied.
Must have the same rank as `intended_shape`.
Returns:
sanitized_slices: Python `list` of slice objects.
"""
sanitized_slices = []
idx = 0
for slc in slices:
if slc is Ellipsis: # Switch over to negative indexing.
if idx < 0:
raise ValueError('Found multiple `...` in slices {}'.format(slices))
num_remaining_non_newaxis_slices = sum(
s is not array_ops.newaxis for s in slices[
slices.index(Ellipsis) + 1:])
idx = -num_remaining_non_newaxis_slices
elif slc is array_ops.newaxis:
pass
else:
is_broadcast = intended_shape[idx] > deficient_shape[idx]
if isinstance(slc, slice):
# Slices are denoted by start:stop:step.
start, stop, step = slc.start, slc.stop, slc.step
if start is not None:
start = _prefer_static_where(is_broadcast, 0, start)
if stop is not None:
stop = _prefer_static_where(is_broadcast, 1, stop)
if step is not None:
step = _prefer_static_where(is_broadcast, 1, step)
slc = slice(start, stop, step)
else: # int, or int Tensor, e.g. d[d.batch_shape_tensor()[0] // 2]
slc = _prefer_static_where(is_broadcast, 0, slc)
idx += 1
sanitized_slices.append(slc)
return sanitized_slices
def _slice_single_param(
param, param_ndims_to_matrix_ndims, slices, batch_shape):
"""Slices into the batch shape of a single parameter.
Args:
param: The original parameter to slice; either a `Tensor` or an object
with batch shape (LinearOperator).
param_ndims_to_matrix_ndims: `int` number of right-most dimensions used for
inferring matrix shape of the `LinearOperator`. For non-Tensor
parameters, this is the number of this param's batch dimensions used by
the matrix shape of the parent object.
slices: iterable of slices received by `__getitem__`.
batch_shape: The parameterized object's batch shape `Tensor`.
Returns:
new_param: Instance of the same type as `param`, batch-sliced according to
`slices`.
"""
# Broadcast the parammeter to have full batch rank.
param = _broadcast_parameter_with_batch_shape(
param, param_ndims_to_matrix_ndims, array_ops.ones_like(batch_shape))
if hasattr(param, 'batch_shape_tensor'):
param_batch_shape = param.batch_shape_tensor()
else:
param_batch_shape = array_ops.shape(param)
# Truncate by param_ndims_to_matrix_ndims
param_batch_rank = array_ops.size(param_batch_shape)
param_batch_shape = param_batch_shape[
:(param_batch_rank - param_ndims_to_matrix_ndims)]
# At this point the param should have full batch rank, *unless* it's an
# atomic object like `tfb.Identity()` incapable of having any batch rank.
if (tensor_util.constant_value(array_ops.size(batch_shape)) != 0 and
tensor_util.constant_value(array_ops.size(param_batch_shape)) == 0):
return param
param_slices = _sanitize_slices(
slices, intended_shape=batch_shape, deficient_shape=param_batch_shape)
# Extend `param_slices` (which represents slicing into the
# parameter's batch shape) with the parameter's event ndims. For example, if
# `params_ndims == 1`, then `[i, ..., j]` would become `[i, ..., j, :]`.
if param_ndims_to_matrix_ndims > 0:
if Ellipsis not in [
slc for slc in slices if not tensor_util.is_tensor(slc)]:
param_slices.append(Ellipsis)
param_slices += [slice(None)] * param_ndims_to_matrix_ndims
return param.__getitem__(tuple(param_slices))
def batch_slice(linop, params_overrides, slices):
"""Slices `linop` along its batch dimensions.
Args:
linop: A `LinearOperator` instance.
params_overrides: A `dict` of parameter overrides.
slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`
thereof. (e.g. the argument of a `__getitem__` method).
Returns:
new_linop: A batch-sliced `LinearOperator`.
"""
if not isinstance(slices, collections.abc.Sequence):
slices = (slices,)
if len(slices) == 1 and slices[0] is Ellipsis:
override_dict = {}
else:
batch_shape = linop.batch_shape_tensor()
override_dict = {}
for param_name, param_ndims_to_matrix_ndims in linop._experimental_parameter_ndims_to_matrix_ndims.items(): # pylint:disable=protected-access,line-too-long
param = getattr(linop, param_name)
# These represent optional `Tensor` parameters.
if param is not None:
override_dict[param_name] = nest.map_structure_up_to(
param, functools.partial(
_slice_single_param, slices=slices, batch_shape=batch_shape),
param, param_ndims_to_matrix_ndims)
override_dict.update(params_overrides)
parameters = dict(linop.parameters, **override_dict)
return type(linop)(**parameters)
tensorflow/python/ops/linalg/sparse/BUILD
+93
View File
@@ -0,0 +1,93 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
# Description: Sparse CSR support for TensorFlow.
load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
tf_gen_op_wrapper_py(
name = "gen_sparse_csr_matrix_ops",
out = "gen_sparse_csr_matrix_ops.py",
api_def_srcs = ["//tensorflow/core/api_def:base_api_def"],
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
py_lib_rule = py_library,
visibility = ["//visibility:private"],
deps = ["//tensorflow/core:sparse_csr_matrix_ops_op_lib"],
)
py_library(
name = "sparse_py",
srcs = ["sparse.py"],
strict_deps = True,
deps = [
":conjugate_gradient",
":sparse_csr_matrix_grad",
":sparse_csr_matrix_ops",
],
)
py_library(
name = "sparse_csr_matrix_grad",
srcs = ["sparse_csr_matrix_grad.py"],
strict_deps = True,
deps = [
":sparse_csr_matrix_ops",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
],
)
py_library(
name = "__init__",
srcs = ["__init__.py"],
strict_deps = True,
)
pytype_strict_library(
name = "conjugate_gradient",
srcs = ["conjugate_gradient.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/linalg:linalg_impl",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "sparse_csr_matrix_ops",
srcs = ["sparse_csr_matrix_ops.py"],
strict_deps = True,
deps = [
":gen_sparse_csr_matrix_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:cpp_shape_inference_proto_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
],
)
tensorflow/python/ops/linalg/sparse/__init__.py
View File
tensorflow/python/ops/linalg/sparse/conjugate_gradient.py
+138
View File
@@ -0,0 +1,138 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Preconditioned Conjugate Gradient."""
import collections
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.ops.linalg import linalg_impl as linalg
from tensorflow.python.util import dispatch
from tensorflow.python.util.tf_export import tf_export
@tf_export('linalg.experimental.conjugate_gradient')
@dispatch.add_dispatch_support
def conjugate_gradient(operator,
rhs,
preconditioner=None,
x=None,
tol=1e-5,
max_iter=20,
name='conjugate_gradient'):
r"""Conjugate gradient solver.
Solves a linear system of equations `A*x = rhs` for self-adjoint, positive
definite matrix `A` and right-hand side vector `rhs`, using an iterative,
matrix-free algorithm where the action of the matrix A is represented by
`operator`. The iteration terminates when either the number of iterations
exceeds `max_iter` or when the residual norm has been reduced to `tol`
times its initial value, i.e. \\(||rhs - A x_k|| <= tol ||rhs||\\).
Args:
operator: A `LinearOperator` that is self-adjoint and positive definite.
rhs: A possibly batched vector of shape `[..., N]` containing the right-hand
size vector.
preconditioner: A `LinearOperator` that approximates the inverse of `A`.
An efficient preconditioner could dramatically improve the rate of
convergence. If `preconditioner` represents matrix `M`(`M` approximates
`A^{-1}`), the algorithm uses `preconditioner.apply(x)` to estimate
`A^{-1}x`. For this to be useful, the cost of applying `M` should be
much lower than computing `A^{-1}` directly.
x: A possibly batched vector of shape `[..., N]` containing the initial
guess for the solution.
tol: A float scalar convergence tolerance.
max_iter: An integer giving the maximum number of iterations.
name: A name scope for the operation.
Returns:
output: A namedtuple representing the final state with fields:
- i: A scalar `int32` `Tensor`. Number of iterations executed.
- x: A rank-1 `Tensor` of shape `[..., N]` containing the computed
solution.
- r: A rank-1 `Tensor` of shape `[.., M]` containing the residual vector.
- p: A rank-1 `Tensor` of shape `[..., N]`. `A`-conjugate basis vector.
- gamma: \\(r \dot M \dot r\\), equivalent to \\(||r||_2^2\\) when
`preconditioner=None`.
"""
if not (operator.is_self_adjoint and operator.is_positive_definite):
raise ValueError('Expected a self-adjoint, positive definite operator.')
cg_state = collections.namedtuple('CGState', ['i', 'x', 'r', 'p', 'gamma'])
def stopping_criterion(i, state):
return math_ops.logical_and(
i < max_iter,
math_ops.reduce_any(linalg.norm(state.r, axis=-1) > tol))
def dot(x, y):
return array_ops.squeeze(
math_ops.matvec(
x[..., array_ops.newaxis],
y, adjoint_a=True), axis=-1)
def cg_step(i, state): # pylint: disable=missing-docstring
z = math_ops.matvec(operator, state.p)
alpha = state.gamma / dot(state.p, z)
x = state.x + alpha[..., array_ops.newaxis] * state.p
r = state.r - alpha[..., array_ops.newaxis] * z
if preconditioner is None:
q = r
else:
q = preconditioner.matvec(r)
gamma = dot(r, q)
beta = gamma / state.gamma
p = q + beta[..., array_ops.newaxis] * state.p
return i + 1, cg_state(i + 1, x, r, p, gamma)
# We now broadcast initial shapes so that we have fixed shapes per iteration.
with ops.name_scope(name):
broadcast_shape = array_ops.broadcast_dynamic_shape(
array_ops.shape(rhs)[:-1],
operator.batch_shape_tensor())
if preconditioner is not None:
broadcast_shape = array_ops.broadcast_dynamic_shape(
broadcast_shape,
preconditioner.batch_shape_tensor()
)
broadcast_rhs_shape = array_ops.concat([
broadcast_shape, [array_ops.shape(rhs)[-1]]], axis=-1)
r0 = array_ops.broadcast_to(rhs, broadcast_rhs_shape)
tol *= linalg.norm(r0, axis=-1)
if x is None:
x = array_ops.zeros(
broadcast_rhs_shape, dtype=rhs.dtype.base_dtype)
else:
r0 = rhs - math_ops.matvec(operator, x)
if preconditioner is None:
p0 = r0
else:
p0 = math_ops.matvec(preconditioner, r0)
gamma0 = dot(r0, p0)
i = constant_op.constant(0, dtype=dtypes.int32)
state = cg_state(i=i, x=x, r=r0, p=p0, gamma=gamma0)
_, state = while_loop.while_loop(stopping_criterion, cg_step, [i, state])
return cg_state(
state.i,
x=state.x,
r=state.r,
p=state.p,
gamma=state.gamma)
tensorflow/python/ops/linalg/sparse/sparse.py
+26
View File
@@ -0,0 +1,26 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Public API for tf.linalg.sparse namespace."""
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.ops.linalg.sparse.conjugate_gradient import conjugate_gradient
from tensorflow.python.ops.linalg.sparse.sparse_csr_matrix_grad import *
from tensorflow.python.ops.linalg.sparse.sparse_csr_matrix_ops import *
# pylint: enable=wildcard-import
__all__ = [
'conjugate_gradient'
]
tensorflow/python/ops/linalg/sparse/sparse_csr_matrix_grad.py
+364
View File
@@ -0,0 +1,364 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""CSR Sparse Matrix Gradients."""
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops
@ops.RegisterGradient("DenseToCSRSparseMatrix")
def _DenseToCSRSparseMatrixGrad(op: ops.Operation, grad):
"""Gradient for dense_to_csr_sparse_matrix op."""
grad_values = (
sparse_csr_matrix_ops.csr_sparse_matrix_to_dense(
grad, type=op.get_attr("T")))
# inputs to fw op were: params, indices.
return (grad_values, None)
@ops.RegisterGradient("CSRSparseMatrixToDense")
def _CSRSparseMatrixToDenseGrad(op: ops.Operation, grad):
"""Gradient for csr_sparse_matrix_to_dense op."""
coo_sparse_tensor = sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
op.inputs[0], type=grad.dtype)
return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
indices=coo_sparse_tensor.indices,
values=array_ops.gather_nd(grad, coo_sparse_tensor.indices),
dense_shape=grad.shape)
@ops.RegisterGradient("SparseTensorToCSRSparseMatrix")
def _SparseTensorToCSRSparseMatrixGrad(op: ops.Operation, grad):
"""Gradient for sparse_tensor_to_csr_sparse_matrix op."""
grad_values = sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
grad, type=op.get_attr("T")).values
return (None, grad_values, None)
@ops.RegisterGradient("CSRSparseMatrixToSparseTensor")
def _CSRSparseMatrixToSparseTensorGrad(op: ops.Operation, *grads):
"""Gradient for csr_sparse_matrix_to_sparse_tensor op."""
return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
indices=op.outputs[0], values=grads[1], dense_shape=op.outputs[2])
ops.NotDifferentiable("SparseMatrixNNZ")
ops.NotDifferentiable("SparseMatrixZeros")
def _PruneSparseTensor(unpruned, pruned_pattern):
"""Helper function to prune COO sparse tensor.
Given two sparse tensors 'unpruned' and 'pruned_pattern', generates another
sparse tensor with indices and values fron 'unpruned' only if its indices also
occur in pruned_pattern.
Args:
unpruned: COO matrix with unpruned indices
pruned_pattern: COO matrix with pruned pattern.
TODO(tabakg): This is far from optimal. Consider a C++ implementation.
Returns:
Indices, values, and dense_shape of the pruned matrix.
"""
pruned_indices = sparse_ops.sparse_reshape(
pruned_pattern, shape=(-1,)).indices[..., 0]
unpruned_indices = sparse_ops.sparse_reshape(
unpruned, shape=(-1,)).indices[..., 0]
best_match = array_ops.searchsorted(unpruned_indices, pruned_indices)
keep_indices = array_ops.gather(
best_match,
array_ops.where(
math_ops.equal(
array_ops.gather(unpruned_indices, best_match), pruned_indices)))
return (array_ops.gather_nd(unpruned.indices, keep_indices),
array_ops.gather_nd(unpruned.values,
keep_indices), pruned_pattern.dense_shape)
def _PruneCSRMatrix(unpruned, pruned_pattern):
"""TODO(tabakg): Consider re-writing in C++."""
_, dtype = sparse_csr_matrix_ops.dense_shape_and_type(pruned_pattern)
coo_unpruned = sparse_tensor.SparseTensor(
*sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
unpruned, type=dtype))
coo_pruned_pattern = sparse_tensor.SparseTensor(
*sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
pruned_pattern, type=dtype))
return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
*_PruneSparseTensor(coo_unpruned, coo_pruned_pattern))
@ops.RegisterGradient("SparseMatrixAdd")
def _SparseMatrixAddGrad(op: ops.Operation, grad):
"""Gradient for sparse_matrix_add op."""
# input to sparse_matrix_add is (a, b, alpha, beta)
# with a, b CSR and alpha beta scalars.
# output is: alpha * a + beta * b
# d(a*A + b*B)/dA . grad = a * grad
# May have gotten the transposes wrong below.
# d(a*A + b*B)/da . grad = tr(A' . grad)
# For now, only implement gradients w.r.t. A and B.
# TODO(ebrevdo): Implement reduce_sum for SparseMatrix so that we
# can implement gradients w.r.t. a and b.
(a_csr, b_csr, alpha, beta) = op.inputs
return (sparse_csr_matrix_ops.sparse_matrix_mul(
_PruneCSRMatrix(grad, a_csr), alpha),
sparse_csr_matrix_ops.sparse_matrix_mul(
_PruneCSRMatrix(grad, b_csr), beta), None, None)
def _PrunedDenseMatrixMultiplication(a,
b,
indices,
transpose_a=False,
adjoint_a=False,
transpose_b=False,
adjoint_b=False):
"""Multiplies two dense matrices at selected indices.
The two inputs `a` and `b` must have matching rank (2 or 3). If using rank 3,
the first rank is used for the batch number. The last two dimensions should
also be compatible for matrix multiplication.
TODO(tabakg): Consider C++ implementation. There is also a more efficient way
to handle transposes here.
Args:
a: The left dense matrix (or batched matrices).
b: The right dense matrix (or batched matrices).
indices: The selected output indices where values should be produced. Other
indices will be pruned (not computed in the first place). Indices are
specified as a tensor of shape (length, rank), where length is the number
of entries and rank is the rank of the dense inputs (2 or 3).
transpose_a: Whether to transpose a.
adjoint_a: Whether to take the conjugate transpose of a.
transpose_b: Whether to transpose b.
adjoint_b: Whether to take the conjugate transpose of b.
Returns:
A CSR matrix.
"""
transpose_a = transpose_a or adjoint_a
transpose_b = transpose_b or adjoint_b
a = math_ops.conj(a) if adjoint_a else a
b = math_ops.conj(b) if adjoint_b else b
rank = len(a.shape)
dense_shape = (a.shape[-1] if transpose_a else a.shape[-2],
b.shape[-2] if transpose_b else b.shape[-1])
if rank == 2:
rows = indices[:, 0]
cols = indices[:, 1]
transpose = array_ops.transpose
gather_op = array_ops.gather
elif rank == 3:
dense_shape = (a.shape[0],) + dense_shape
rows = indices[:, :2]
cols = array_ops_stack.stack([indices[:, 0], indices[:, 2]], axis=1)
transpose = lambda x: array_ops.transpose(x, perm=[0, 2, 1])
gather_op = array_ops.gather_nd
a_rows = gather_op(transpose(a) if transpose_a else a, indices=rows)
b_cols = gather_op(b if transpose_b else transpose(b), indices=cols)
values = math_ops.reduce_sum(a_rows * b_cols, axis=1)
return sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
indices=indices, values=values, dense_shape=dense_shape)
@ops.RegisterGradient("SparseMatrixTranspose")
def _SparseMatrixTransposeGrad(op: ops.Operation, grad):
"""Gradient for sparse_matrix_transpose op."""
return sparse_csr_matrix_ops.sparse_matrix_transpose(
grad, type=op.get_attr("type"), conjugate=op.get_attr("conjugate"))
@ops.RegisterGradient("SparseMatrixSoftmax")
def _SparseMatrixSoftmaxGrad(op: ops.Operation, grad_softmax):
"""Gradient for sparse_matrix_softmax op."""
softmax = op.outputs[0]
return sparse_csr_matrix_ops.sparse_matrix_softmax_grad(
softmax, grad_softmax, type=op.get_attr("type"))
@ops.RegisterGradient("SparseMatrixMatMul")
def _SparseMatrixMatMulGrad(op: ops.Operation, grad):
"""Gradient for sparse_matrix_mat_mul op."""
# input to sparse_matrix_mat_mul is (A, B) with CSR A and dense B.
# Output is dense:
# C = opA(A) . opB(B) if transpose_output = false
# C = (opA(A) . opB(B))' = opB(B)' . opA(A)' if transpose_output = true.
# where opA = transpose if transpose_a = True else identity
# and opB = transpose if transpose_b = True else identity
t_a = op.get_attr("transpose_a")
t_b = op.get_attr("transpose_b")
adj_a = op.get_attr("adjoint_a")
adj_b = op.get_attr("adjoint_b")
transpose_output = op.get_attr("transpose_output")
conjugate_output = op.get_attr("conjugate_output")
a = op.inputs[0] # sparse matrix
b = op.inputs[1] # dense matrix
conj = math_ops.conj
sparse_matmul = sparse_csr_matrix_ops.sparse_matrix_mat_mul
def matmul(x, y, **kwargs): # pylint: disable=invalid-name
return _PrunedDenseMatrixMultiplication(
x,
y,
indices=sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
a, type=x.dtype).indices,
**kwargs)
if conjugate_output:
grad = conj(grad)
if not transpose_output:
# C = opA(A) . opB(B)
if not adj_a and not adj_b:
a = conj(a)
b = conj(b)
if not t_a:
grad_a = matmul(grad, b, transpose_b=not t_b)
else:
grad_a = matmul(b, grad, transpose_a=t_b, transpose_b=True)
grad_b = sparse_matmul(a, grad, transpose_a=not t_a, transpose_output=t_b)
elif not t_a and not t_b:
if not adj_a:
grad_a = matmul(grad, b, adjoint_b=not adj_b)
else:
grad_a = matmul(b, grad, adjoint_a=adj_b, adjoint_b=True)
grad_b = sparse_matmul(
a,
grad,
adjoint_a=not adj_a,
transpose_output=adj_b,
conjugate_output=adj_b)
elif adj_a and t_b:
grad_a = matmul(b, grad, transpose_a=True, adjoint_b=True)
grad_b = sparse_matmul(a, grad, transpose_output=True)
elif t_a and adj_b:
grad_a = matmul(b, grad, transpose_a=True, transpose_b=True)
grad_b = sparse_matmul(
conj(a), grad, transpose_output=True, conjugate_output=True)
else:
# C = (opA(A) . opB(B))' = opB(B)' . opA(A)'
if not adj_a and not adj_b:
a = conj(a)
b = conj(b)
if not t_a:
grad_a = matmul(grad, b, transpose_a=True, transpose_b=not t_b)
else:
grad_a = matmul(b, grad, transpose_a=t_b)
grad_b = sparse_matmul(
a, grad, transpose_a=not t_a, transpose_b=True, transpose_output=t_b)
elif not t_a and not t_b:
if not adj_a:
grad_a = matmul(grad, b, transpose_a=True, adjoint_b=not adj_b)
else:
grad_a = matmul(b, conj(grad), adjoint_a=adj_b)
grad_b = sparse_matmul(
a,
grad,
adjoint_a=not adj_a,
transpose_b=True,
transpose_output=adj_b,
conjugate_output=adj_b)
elif adj_a and t_b:
grad_a = matmul(b, conj(grad), transpose_a=True)
grad_b = sparse_matmul(a, grad, transpose_b=True, transpose_output=True)
elif t_a and adj_b:
grad_a = matmul(b, grad, transpose_a=True)
grad_b = sparse_matmul(a, grad, adjoint_b=True, transpose_output=True)
return (grad_a, grad_b)
@ops.RegisterGradient("SparseMatrixSparseMatMul")
def _SparseMatrixSparseMatMulGrad(op: ops.Operation, grad):
"""Gradient for sparse_matrix_sparse_mat_mul op."""
t_a = op.get_attr("transpose_a")
t_b = op.get_attr("transpose_b")
adj_a = op.get_attr("adjoint_a")
adj_b = op.get_attr("adjoint_b")
dtype = op.get_attr("type")
# input to sparse_matrix_sparse_mat_mul is (A, B) with CSR A and B.
# Output is CSR:
# C = opA(A) . opB(B)
# where opA = transpose if transpose_a = True else identity
# and opB = transpose if transpose_b = True else identity
a = op.inputs[0]
b = op.inputs[1]
conj = math_ops.conj
matmul = sparse_csr_matrix_ops.sparse_matrix_sparse_mat_mul
if not t_a and not t_b:
if not adj_a:
if not adj_b:
grad_a = matmul(grad, b, adjoint_b=True, type=dtype)
grad_b = matmul(a, grad, adjoint_a=True, type=dtype)
else:
grad_a = matmul(grad, b, type=dtype)
grad_b = matmul(grad, a, adjoint_a=True, type=dtype)
else:
if not adj_b:
grad_a = matmul(b, grad, adjoint_b=True, type=dtype)
grad_b = matmul(a, grad, type=dtype)
else:
grad_a = matmul(b, grad, adjoint_a=True, adjoint_b=True, type=dtype)
grad_b = matmul(grad, a, adjoint_a=True, adjoint_b=True, type=dtype)
elif not adj_a and not adj_b:
if not t_a and t_b:
grad_a = matmul(grad, conj(b), type=dtype)
grad_b = matmul(grad, conj(a), transpose_a=True, type=dtype)
elif t_a and not t_b:
grad_a = matmul(conj(b), grad, transpose_b=True, type=dtype)
grad_b = matmul(conj(a), grad, type=dtype)
else:
grad_a = matmul(b, grad, adjoint_a=True, transpose_b=True, type=dtype)
grad_b = matmul(grad, a, transpose_a=True, adjoint_b=True, type=dtype)
elif adj_a and t_b:
grad_a = matmul(b, grad, transpose_a=True, adjoint_b=True, type=dtype)
grad_b = matmul(grad, a, transpose_a=True, transpose_b=True, type=dtype)
elif t_a and adj_b:
grad_a = matmul(b, grad, transpose_a=True, transpose_b=True, type=dtype)
grad_b = matmul(grad, a, adjoint_a=True, transpose_b=True, type=dtype)
# TODO(tabakg): There should be a C++ function for sparse-sparse
# multiplication with pre-determined indices, instead of pruning after the
# multiplication.
return (_PruneCSRMatrix(grad_a, a), _PruneCSRMatrix(grad_b, b))
@ops.RegisterGradient("SparseMatrixMul")
def _SparseMatrixMulGrad(op: ops.Operation, grad):
"""Gradient for sparse_matrix_mul op."""
# input to sparse_matrix_mul is (A, B) with CSR A and dense B.
# Output is CSR:
# C = A .* B
del op
del grad
raise NotImplementedError
tensorflow/python/ops/linalg/sparse/sparse_csr_matrix_ops.py
+376
View File
@@ -0,0 +1,376 @@
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""CSR Sparse Matrix Operations."""
import abc
import collections
# pylint: disable=g-direct-tensorflow-import, wildcard-import
from tensorflow.python.eager import context
from tensorflow.python.framework import cpp_shape_inference_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops.linalg.sparse import gen_sparse_csr_matrix_ops as sm_ops
from tensorflow.python.ops.linalg.sparse.gen_sparse_csr_matrix_ops import *
__all__ = [
"SparseMatrix",
"CSRSparseMatrix",
"matmul",
"dense_shape_and_type",
]
# pylint: disable=invalid-name
__all__ += [_x for _x in dir(sm_ops) if not _x.startswith("_")]
class DenseShapeAndType(
collections.namedtuple("DenseShapeAndType", ("shape", "dtype"))):
pass
def _get_handle_data(tensor):
return resource_variable_ops.get_eager_safe_handle_data(tensor)
def _create_handle_data_proto(shape_proto, dtype_enum):
"""Create handle data based on shape and dtype protos."""
variant_shape_and_type_data = \
cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
variant_shape_and_type_data.is_set = True
# NOTE(ebrevdo): shape_and_type lacks append() in some versions of protobuf.
variant_shape_and_type_data.shape_and_type.extend([
cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
shape=shape_proto, dtype=dtype_enum)
])
return variant_shape_and_type_data
def _make_handle_data(tensor):
"""Create handle data based on tensor shape and dtype."""
return _create_handle_data_proto(tensor.shape.as_proto(),
tensor.dtype.as_datatype_enum)
def get_shape_and_type(matrix):
"""Return matrix's shape and type if available."""
handle_data = getattr(matrix, "_handle_data", None)
if handle_data is None:
return None
if len(handle_data.shape_and_type) != 1:
raise ValueError(
"shape_and_type array in _handle_data must have length one, but saw: %d"
% len(handle_data.shape_and_type))
return handle_data.shape_and_type[0]
def dense_shape_and_type(matrix):
"""Get dense shape and dtype of the tf.Tensor containing the matrix.
Args:
matrix: A `tf.Tensor` of type `tf.variant` storing a sparse matrix.
Returns:
An instance of `ShapeAndType` with properties `shape` (a `tf.TensorShape`)
and `dtype` (a `tf.DType`).
Raises:
TypeError: if `matrix` is not a tensor or its dtype is not variant.
ValueError: if `matrix` lacks static handle data containing the dense
shape and dtype.
"""
if not isinstance(matrix, tensor_lib.Tensor):
raise TypeError("matrix should be a tensor, but saw: %s" % (matrix,))
if matrix.dtype != dtypes.variant:
raise TypeError(
"expected matrix to be type tf.variant, but saw: %s" % (matrix.dtype,))
handle_data = _get_handle_data(matrix)
if not handle_data or not handle_data.is_set:
raise ValueError("matrix has missing handle data: %s" % (matrix,))
if len(handle_data.shape_and_type) != 1:
raise ValueError("len(matrix.handle_data.shape_and_type) != 1: '%s'" %
(handle_data.shape_and_type,))
return DenseShapeAndType(
tensor_shape.TensorShape(handle_data.shape_and_type[0].shape),
dtypes.DType(handle_data.shape_and_type[0].dtype))
def matmul_shape_inference(a, b, c, transpose_a, transpose_b, adjoint_a,
adjoint_b):
"""Helper function for matmul to set the result matrix's handle data."""
c_handle = getattr(c, "_handle_data", None)
a_shape_and_type = get_shape_and_type(a)
b_shape_and_type = get_shape_and_type(b)
if (c_handle is None and a_shape_and_type is not None and
b_shape_and_type is not None):
transpose_a = transpose_a or adjoint_a
transpose_b = transpose_b or adjoint_b
a_shape = a_shape_and_type.shape
b_shape = b_shape_and_type.shape
rank = len(a_shape.dim)
# Creates the output shape.
c_rows = a_shape.dim[rank - (1 if transpose_a else 2)].size
c_cols = b_shape.dim[rank - (2 if transpose_b else 1)].size
c_shape = tensor_shape.TensorShape(a_shape)
c_shape = tensor_shape.TensorShape(c_shape[:rank - 2] + [c_rows, c_cols])
c_handle = _create_handle_data_proto(c_shape.as_proto(),
a_shape_and_type.dtype)
return c_handle
def matmul(a,
b,
transpose_a=False,
transpose_b=False,
adjoint_a=False,
adjoint_b=False,
name=None):
"""Perform a sparse matrix matmul between `a` and `b`.
Performs a contraction between `a` and `b` along the two innermost dimensions.
If both `a` and `b` are instances of `SparseMatrix`, returns a new instance
of `SparseMatrix` (same type as `a`). If one is not an instance of
`SparseMatrix`, returns a dense `Tensor`:
```
c = opA(a) . opB(b)
```
where `opA` (resp. `opB`) is the transpose or hermitian transpose depending
on the values of `transpose_a` (resp. `transpose_b`) and `adjoint_a`
(resp. `adjoint_b`).
Args:
a: `Tensor` or `SparseMatrix`, having rank `2` or `3`.
b: `Tensor` or `SparseMatrix`, having rank `2` or `3`.
transpose_a: Python `bool`.
transpose_b: Python `bool`.
adjoint_a: Python `bool`.
adjoint_b: Python `bool`.
name: Optional name to use when creating ops.
Returns:
A `SparseMatrix` if both `a` and `b` are instances of `SparseMatrix`,
otherwise a dense `Tensor`.
"""
if not isinstance(a, SparseMatrix) and not isinstance(b, SparseMatrix):
return math_ops.matmul(
a,
b,
transpose_a=transpose_a,
transpose_b=transpose_b,
adjoint_a=adjoint_a,
adjoint_b=adjoint_b,
name=name)
# pylint: disable=protected-access
a_matrix = a._matrix if isinstance(a, SparseMatrix) else a
b_matrix = b._matrix if isinstance(b, SparseMatrix) else b
with ops.name_scope(name, "SparseMatrixMatMul", [a_matrix, b_matrix]):
if isinstance(a, SparseMatrix) and isinstance(b, SparseMatrix):
if not (isinstance(a, type(b)) or isinstance(b, type(a))):
raise TypeError("SparseMatrix types don't inherit from each other: "
"%s and %s" % (type(a), type(b)))
c = sm_ops.sparse_matrix_sparse_mat_mul(
a_matrix,
b_matrix,
transpose_a=transpose_a,
transpose_b=transpose_b,
adjoint_a=adjoint_a,
adjoint_b=adjoint_b,
type=a.dtype)
# In eager mode, shape inference functions are not called, and the output
# shape is not set. We have to infer the output shape here.
# TODO(penporn): Set this from the C++ kernel instead.
c_handle = matmul_shape_inference(a_matrix, b_matrix, c, transpose_a,
transpose_b, adjoint_a, adjoint_b)
return a._from_matrix(c, handle_data=c_handle)
elif isinstance(a, SparseMatrix):
return sm_ops.sparse_matrix_mat_mul(
a_matrix,
b,
transpose_a=transpose_a,
transpose_b=transpose_b,
adjoint_a=adjoint_a,
adjoint_b=adjoint_b)
else:
# opA(A) . opB(B) = t(nopB(B) . nopA(A))
if not adjoint_a and not adjoint_b:
return sm_ops.sparse_matrix_mat_mul(
b_matrix,
a,
transpose_a=not transpose_b,
transpose_b=not transpose_a,
transpose_output=True)
elif not transpose_a and not transpose_b:
return sm_ops.sparse_matrix_mat_mul(
b_matrix,
a,
adjoint_a=not adjoint_b,
adjoint_b=not adjoint_a,
transpose_output=True,
conjugate_output=True)
else:
return sm_ops.sparse_matrix_mat_mul(
b_matrix,
math_ops.conj(a),
transpose_output=True,
conjugate_output=adjoint_b)
class SparseMatrix(metaclass=abc.ABCMeta):
"""Abstract class for sparse matrix types."""
@abc.abstractmethod
def __init__(self):
self._eager_mode = context.executing_eagerly()
@abc.abstractproperty
def _matrix(self):
pass
@abc.abstractmethod
def _from_matrix(self, matrix, handle_data=None):
pass
@abc.abstractmethod
def to_dense(self):
pass
@abc.abstractmethod
def to_sparse_tensor(self):
pass
@property
def graph(self):
return self._matrix.graph
@property
def shape(self):
return dense_shape_and_type(self._matrix).shape
@property
def dtype(self):
return dense_shape_and_type(self._matrix).dtype
@property
def eager_handle_data(self):
"""Return the matrix's handle data iff in eager mode."""
return _get_handle_data(self._matrix) if self._eager_mode else None
def conj(self):
return self._from_matrix(
math_ops.conj(self._matrix), self.eager_handle_data)
def hermitian_transpose(self):
"""Return the hermitian transpose of the matrix."""
return self._from_matrix(
sm_ops.sparse_matrix_transpose(
self._matrix, conjugate=True, type=self.dtype),
self.eager_handle_data)
def nnz(self):
"""Number of stored values, including explicit zeros."""
return sm_ops.sparse_matrix_nnz(self._matrix)
nonzero = nnz
def sorted_indices(self):
# TODO(ebrevdo): A more efficient implementation?
return self.to_sparse_tensor().indices
def transpose(self):
return self._from_matrix(
sm_ops.sparse_matrix_transpose(self._matrix, type=self.dtype),
self.eager_handle_data)
class CSRSparseMatrix(SparseMatrix):
"""(Optionally batched) CSR Sparse Matrix."""
def __init__(self, value, indices=None, name=None):
"""Construct a CSRSparseMatrix from a dense matrix or SparseTensor.
Args:
value: A dense `2D` or `3D` Tensor or `SparseTensor`.
indices: The nonzero indices of `value`
(if `value` is not a `SparseTensor`).
name: Optional op name.
Raises:
ValueError: if `value` is a `SparseTensor` and `indices` is not `None`.
"""
del name # Unused.
super(CSRSparseMatrix, self).__init__()
if isinstance(value, sparse_tensor.SparseTensor):
if indices is not None:
raise ValueError("indices must be None if value is a SparseTensor.")
self._dtype = value.dtype
self._csr_matrix = sm_ops.sparse_tensor_to_csr_sparse_matrix(
indices=value.indices,
values=value.values,
dense_shape=value.dense_shape)
else:
value = ops.convert_to_tensor(value)
self._dtype = value.dtype
if indices is not None:
indices = ops.convert_to_tensor(indices, dtype=dtypes.int64)
else:
indices = array_ops.stop_gradient(array_ops.where(value))
self._csr_matrix = sm_ops.dense_to_csr_sparse_matrix(value, indices)
# Eager mode doesn't call shape inference functions, so we have to set the
# shape and dtype handle data directly.
if self._eager_mode:
# pylint: disable=protected-access
self._csr_matrix._handle_data = _make_handle_data(value)
# pylint: enable=protected-access
@property
def _matrix(self):
return self._csr_matrix
def _from_matrix(self, matrix, handle_data=None):
assert (
isinstance(matrix, tensor_lib.Tensor) and matrix.dtype == dtypes.variant
)
ret = type(self).__new__(type(self))
# pylint: disable=protected-access
ret._dtype = self._dtype
if self._eager_mode:
if matrix._handle_data is None:
matrix._handle_data = handle_data
assert matrix._handle_data is not None
ret._csr_matrix = matrix
# pylint: enable=protected-access
return ret
def to_dense(self):
return sm_ops.csr_sparse_matrix_to_dense(self._matrix, type=self.dtype)
def to_sparse_tensor(self):
r = sm_ops.csr_sparse_matrix_to_sparse_tensor(self._matrix, type=self.dtype)
return sparse_tensor.SparseTensor(
indices=r.indices, values=r.values, dense_shape=r.dense_shape)