xref: /aosp_15_r20/external/tensorflow/tensorflow/python/ops/linalg_ops.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Operations for linear algebra."""
16
17import numpy as np
18
19from tensorflow.python.framework import dtypes
20from tensorflow.python.framework import ops
21from tensorflow.python.ops import array_ops
22from tensorflow.python.ops import control_flow_ops
23from tensorflow.python.ops import gen_array_ops
24from tensorflow.python.ops import gen_linalg_ops
25from tensorflow.python.ops import linalg_ops_impl
26from tensorflow.python.ops import map_fn
27from tensorflow.python.ops import math_ops
28# pylint: disable=wildcard-import
29from tensorflow.python.ops.gen_linalg_ops import *
30# pylint: enable=wildcard-import
31from tensorflow.python.util import deprecation
32from tensorflow.python.util import dispatch
33from tensorflow.python.util.tf_export import tf_export
34
35# Names below are lower_case.
36# pylint: disable=invalid-name
37
38
39def _RegularizedGramianCholesky(matrix, l2_regularizer, first_kind):
40  r"""Computes Cholesky factorization of regularized gramian matrix.
41
42  Below we will use the following notation for each pair of matrix and
43  right-hand sides in the batch:
44
45  `matrix`=\\(A \in \Re^{m \times n}\\),
46  `output`=\\(C  \in \Re^{\min(m, n) \times \min(m,n)}\\),
47  `l2_regularizer`=\\(\lambda\\).
48
49  If `first_kind` is True, returns the Cholesky factorization \\(L\\) such that
50  \\(L L^H =  A^H A + \lambda I\\).
51  If `first_kind` is False, returns the Cholesky factorization \\(L\\) such that
52  \\(L L^H =  A A^H + \lambda I\\).
53
54  Args:
55    matrix: `Tensor` of shape `[..., M, N]`.
56    l2_regularizer: 0-D `double` `Tensor`. Ignored if `fast=False`.
57    first_kind: bool. Controls what gramian matrix to factor.
58  Returns:
59    output: `Tensor` of shape `[..., min(M,N), min(M,N)]` whose inner-most 2
60      dimensions contain the Cholesky factors \\(L\\) described above.
61  """
62
63  gramian = math_ops.matmul(
64      matrix, matrix, adjoint_a=first_kind, adjoint_b=not first_kind)
65  if isinstance(l2_regularizer, ops.Tensor) or l2_regularizer != 0:
66    matrix_shape = array_ops.shape(matrix)
67    batch_shape = matrix_shape[:-2]
68    if first_kind:
69      small_dim = matrix_shape[-1]
70    else:
71      small_dim = matrix_shape[-2]
72    identity = eye(small_dim, batch_shape=batch_shape, dtype=matrix.dtype)
73    small_dim_static = matrix.shape[-1 if first_kind else -2]
74    identity.set_shape(
75        matrix.shape[:-2].concatenate([small_dim_static, small_dim_static]))
76    gramian += l2_regularizer * identity
77  return gen_linalg_ops.cholesky(gramian)
78
79
80@tf_export(
81    'linalg.triangular_solve',
82    v1=['linalg.triangular_solve', 'matrix_triangular_solve'])
83@dispatch.add_dispatch_support
84def matrix_triangular_solve(matrix, rhs, lower=True, adjoint=False, name=None):
85  """Solve systems of linear equations with upper or lower triangular matrices.
86
87  `matrix` is a tensor of shape `[..., M, M]` whose inner-most 2 dimensions form
88  square matrices. If `lower` is `True` then the strictly upper triangular part
89  of each inner-most matrix is assumed to be zero and not accessed. If `lower`
90  is `False` then the strictly lower triangular part of each inner-most matrix
91  is assumed to be zero and not accessed. `rhs` is a tensor of shape
92  `[..., M, N]`.
93
94  The output is a tensor of shape `[..., M, N]`. If `adjoint` is `True` then the
95  innermost matrices in output satisfy matrix equations `
96  sum_k matrix[..., i, k] * output[..., k, j] = rhs[..., i, j]`.
97  If `adjoint` is `False` then the
98  innermost matrices in output satisfy matrix equations
99  `sum_k adjoint(matrix[..., i, k]) * output[..., k, j] = rhs[..., i, j]`.
100
101  Example:
102
103  >>> a = tf.constant([[3,  0,  0,  0],
104  ...   [2,  1,  0,  0],
105  ...   [1,  0,  1,  0],
106  ...   [1,  1,  1,  1]], dtype=tf.float32)
107
108  >>> b = tf.constant([[4], [2], [4], [2]], dtype=tf.float32)
109  >>> x = tf.linalg.triangular_solve(a, b, lower=True)
110  >>> x
111  <tf.Tensor: shape=(4, 1), dtype=float32, numpy=
112  array([[ 1.3333334 ],
113         [-0.66666675],
114         [ 2.6666665 ],
115         [-1.3333331 ]], dtype=float32)>
116  >>> tf.matmul(a, x)
117  <tf.Tensor: shape=(4, 1), dtype=float32, numpy=
118  array([[4.],
119         [2.],
120         [4.],
121         [2.]], dtype=float32)>
122
123  Args:
124    matrix: A `Tensor`. Must be one of the following types: `float64`,
125      `float32`, `half`, `complex64`, `complex128`. Shape is `[..., M, M]`.
126    rhs: A `Tensor`. Must have the same type as `matrix`. Shape is `[..., M,
127      N]`.
128    lower: An optional `bool`. Defaults to `True`. Boolean indicating whether
129      the innermost matrices in matrix are lower or upper triangular.
130    adjoint: An optional `bool`. Defaults to `False`. Boolean indicating whether
131      to solve with matrix or its (block-wise) adjoint.
132    name: A name for the operation (optional).
133
134  Returns:
135    A `Tensor`. Has the same type as matrix, and shape is `[..., M, N]`.
136
137  """
138  with ops.name_scope(name, 'triangular_solve', [matrix, rhs]):
139    return gen_linalg_ops.matrix_triangular_solve(
140        matrix, rhs, lower=lower, adjoint=adjoint)
141
142
143@tf_export(
144    'linalg.cholesky_solve', v1=['linalg.cholesky_solve', 'cholesky_solve'])
145@dispatch.add_dispatch_support
146@deprecation.deprecated_endpoints('cholesky_solve')
147def cholesky_solve(chol, rhs, name=None):
148  """Solves systems of linear eqns `A X = RHS`, given Cholesky factorizations.
149
150  Specifically, returns `X` from `A X = RHS`, where `A = L L^T`, `L` is the
151  `chol` arg and `RHS` is the `rhs` arg.
152
153  ```python
154  # Solve 10 separate 2x2 linear systems:
155  A = ... # shape 10 x 2 x 2
156  RHS = ... # shape 10 x 2 x 1
157  chol = tf.linalg.cholesky(A)  # shape 10 x 2 x 2
158  X = tf.linalg.cholesky_solve(chol, RHS)  # shape 10 x 2 x 1
159  # tf.matmul(A, X) ~ RHS
160  X[3, :, 0]  # Solution to the linear system A[3, :, :] x = RHS[3, :, 0]
161
162  # Solve five linear systems (K = 5) for every member of the length 10 batch.
163  A = ... # shape 10 x 2 x 2
164  RHS = ... # shape 10 x 2 x 5
165  ...
166  X[3, :, 2]  # Solution to the linear system A[3, :, :] x = RHS[3, :, 2]
167  ```
168
169  Args:
170    chol:  A `Tensor`.  Must be `float32` or `float64`, shape is `[..., M, M]`.
171      Cholesky factorization of `A`, e.g. `chol = tf.linalg.cholesky(A)`.
172      For that reason, only the lower triangular parts (including the diagonal)
173      of the last two dimensions of `chol` are used.  The strictly upper part is
174      assumed to be zero and not accessed.
175    rhs:  A `Tensor`, same type as `chol`, shape is `[..., M, K]`.
176    name:  A name to give this `Op`.  Defaults to `cholesky_solve`.
177
178  Returns:
179    Solution to `A x = rhs`, shape `[..., M, K]`.
180  """
181  # To solve C C^* x = rhs, we
182  # 1. Solve C y = rhs for y, thus y = C^* x
183  # 2. Solve C^* x = y for x
184  with ops.name_scope(name, 'cholesky_solve', [chol, rhs]):
185    y = gen_linalg_ops.matrix_triangular_solve(
186        chol, rhs, adjoint=False, lower=True)
187    x = gen_linalg_ops.matrix_triangular_solve(
188        chol, y, adjoint=True, lower=True)
189    return x
190
191
192@tf_export('eye', 'linalg.eye')
193@dispatch.add_dispatch_support
194def eye(num_rows,
195        num_columns=None,
196        batch_shape=None,
197        dtype=dtypes.float32,
198        name=None):
199  """Construct an identity matrix, or a batch of matrices.
200
201  See also `tf.ones`, `tf.zeros`, `tf.fill`, `tf.one_hot`.
202
203  ```python
204  # Construct one identity matrix.
205  tf.eye(2)
206  ==> [[1., 0.],
207       [0., 1.]]
208
209  # Construct a batch of 3 identity matrices, each 2 x 2.
210  # batch_identity[i, :, :] is a 2 x 2 identity matrix, i = 0, 1, 2.
211  batch_identity = tf.eye(2, batch_shape=[3])
212
213  # Construct one 2 x 3 "identity" matrix
214  tf.eye(2, num_columns=3)
215  ==> [[ 1.,  0.,  0.],
216       [ 0.,  1.,  0.]]
217  ```
218
219  Args:
220    num_rows: Non-negative `int32` scalar `Tensor` giving the number of rows
221      in each batch matrix.
222    num_columns: Optional non-negative `int32` scalar `Tensor` giving the number
223      of columns in each batch matrix.  Defaults to `num_rows`.
224    batch_shape:  A list or tuple of Python integers or a 1-D `int32` `Tensor`.
225      If provided, the returned `Tensor` will have leading batch dimensions of
226      this shape.
227    dtype:  The type of an element in the resulting `Tensor`
228    name:  A name for this `Op`.  Defaults to "eye".
229
230  Returns:
231    A `Tensor` of shape `batch_shape + [num_rows, num_columns]`
232  """
233  return linalg_ops_impl.eye(num_rows,
234                             num_columns=num_columns,
235                             batch_shape=batch_shape,
236                             dtype=dtype,
237                             name=name)
238
239
240@tf_export('linalg.lstsq', v1=['linalg.lstsq', 'matrix_solve_ls'])
241@dispatch.add_dispatch_support
242@deprecation.deprecated_endpoints('matrix_solve_ls')
243def matrix_solve_ls(matrix, rhs, l2_regularizer=0.0, fast=True, name=None):
244  r"""Solves one or more linear least-squares problems.
245
246  `matrix` is a tensor of shape `[..., M, N]` whose inner-most 2 dimensions
247  form `M`-by-`N` matrices. Rhs is a tensor of shape `[..., M, K]` whose
248  inner-most 2 dimensions form `M`-by-`K` matrices.  The computed output is a
249  `Tensor` of shape `[..., N, K]` whose inner-most 2 dimensions form `M`-by-`K`
250  matrices that solve the equations
251  `matrix[..., :, :] * output[..., :, :] = rhs[..., :, :]` in the least squares
252  sense.
253
254  Below we will use the following notation for each pair of matrix and
255  right-hand sides in the batch:
256
257  `matrix`=\\(A \in \Re^{m \times n}\\),
258  `rhs`=\\(B  \in \Re^{m \times k}\\),
259  `output`=\\(X  \in \Re^{n \times k}\\),
260  `l2_regularizer`=\\(\lambda\\).
261
262  If `fast` is `True`, then the solution is computed by solving the normal
263  equations using Cholesky decomposition. Specifically, if \\(m \ge n\\) then
264  \\(X = (A^T A + \lambda I)^{-1} A^T B\\), which solves the least-squares
265  problem \\(X = \mathrm{argmin}_{Z \in \Re^{n \times k}} ||A Z - B||_F^2 +
266  \lambda ||Z||_F^2\\). If \\(m \lt n\\) then `output` is computed as
267  \\(X = A^T (A A^T + \lambda I)^{-1} B\\), which (for \\(\lambda = 0\\)) is
268  the minimum-norm solution to the under-determined linear system, i.e.
269  \\(X = \mathrm{argmin}_{Z \in \Re^{n \times k}} ||Z||_F^2 \\), subject to
270  \\(A Z = B\\). Notice that the fast path is only numerically stable when
271  \\(A\\) is numerically full rank and has a condition number
272  \\(\mathrm{cond}(A) \lt \frac{1}{\sqrt{\epsilon_{mach}}}\\) or\\(\lambda\\)
273  is sufficiently large.
274
275  If `fast` is `False` an algorithm based on the numerically robust complete
276  orthogonal decomposition is used. This computes the minimum-norm
277  least-squares solution, even when \\(A\\) is rank deficient. This path is
278  typically 6-7 times slower than the fast path. If `fast` is `False` then
279  `l2_regularizer` is ignored.
280
281  Args:
282    matrix: `Tensor` of shape `[..., M, N]`.
283    rhs: `Tensor` of shape `[..., M, K]`.
284    l2_regularizer: 0-D `double` `Tensor`. Ignored if `fast=False`.
285    fast: bool. Defaults to `True`.
286    name: string, optional name of the operation.
287
288  Returns:
289    output: `Tensor` of shape `[..., N, K]` whose inner-most 2 dimensions form
290      `M`-by-`K` matrices that solve the equations
291      `matrix[..., :, :] * output[..., :, :] = rhs[..., :, :]` in the least
292      squares sense.
293
294  Raises:
295    NotImplementedError: linalg.lstsq is currently disabled for complex128
296    and l2_regularizer != 0 due to poor accuracy.
297  """
298
299  # pylint: disable=long-lambda
300  def _use_composite_impl(fast, tensor_shape):
301    """Determines whether to use the composite or specialized CPU kernel.
302
303    When the total size of the tensor is larger than the cache size and the
304    batch size is large compared to the smallest matrix dimension, then the
305    composite implementation is inefficient since it has to read the entire
306    tensor from memory multiple times. In this case we fall back to the
307    original CPU kernel, which does all the computational steps on each
308    matrix separately.
309
310    Only fast mode is supported by the composite impl, so `False` is returned
311    if `fast` is `False`.
312
313    Args:
314      fast: bool indicating if fast mode in the solver was requested.
315      tensor_shape: The shape of the tensor.
316
317    Returns:
318      True if the composite impl should be used. False otherwise.
319    """
320    if fast is False:
321      return False
322    batch_shape = tensor_shape[:-2]
323    matrix_shape = tensor_shape[-2:]
324    if not tensor_shape.is_fully_defined():
325      return True
326    tensor_size = tensor_shape.num_elements() * matrix.dtype.size
327    is_io_bound = batch_shape.num_elements() > np.min(matrix_shape)
328    L2_CACHE_SIZE_GUESSTIMATE = 256000
329    if tensor_size > L2_CACHE_SIZE_GUESSTIMATE and is_io_bound:
330      return False
331    else:
332      return True
333
334  def _overdetermined(matrix, rhs, l2_regularizer):
335    """Computes (A^H*A + l2_regularizer)^{-1} * A^H * rhs."""
336    chol = _RegularizedGramianCholesky(
337        matrix, l2_regularizer=l2_regularizer, first_kind=True)
338    return cholesky_solve(chol, math_ops.matmul(matrix, rhs, adjoint_a=True))
339
340  def _underdetermined(matrix, rhs, l2_regularizer):
341    """Computes A^H * (A*A^H + l2_regularizer)^{-1} * rhs."""
342    chol = _RegularizedGramianCholesky(
343        matrix, l2_regularizer=l2_regularizer, first_kind=False)
344    return math_ops.matmul(matrix, cholesky_solve(chol, rhs), adjoint_a=True)
345
346  def _composite_impl(matrix, rhs, l2_regularizer):
347    """Composite implementation of matrix_solve_ls that supports GPU."""
348    with ops.name_scope(name, 'matrix_solve_ls', [matrix, rhs, l2_regularizer]):
349      matrix_shape = matrix.get_shape()[-2:]
350      if matrix_shape.is_fully_defined():
351        if matrix_shape[-2] >= matrix_shape[-1]:
352          return _overdetermined(matrix, rhs, l2_regularizer)
353        else:
354          return _underdetermined(matrix, rhs, l2_regularizer)
355      else:
356        # We have to defer determining the shape to runtime and use
357        # conditional execution of the appropriate graph.
358        matrix_shape = array_ops.shape(matrix)[-2:]
359        return control_flow_ops.cond(
360            matrix_shape[-2] >= matrix_shape[-1],
361            lambda: _overdetermined(matrix, rhs, l2_regularizer),
362            lambda: _underdetermined(matrix, rhs, l2_regularizer))
363
364  matrix = ops.convert_to_tensor(matrix, name='matrix')
365  if matrix.dtype == dtypes.complex128 and l2_regularizer != 0:
366    # TODO(rmlarsen): Investigate and fix accuracy bug.
367    raise NotImplementedError('matrix_solve_ls is currently disabled for '
368                              'complex128 and l2_regularizer != 0 due to '
369                              'poor accuracy.')
370  tensor_shape = matrix.get_shape()
371  if _use_composite_impl(fast, tensor_shape):
372    return _composite_impl(matrix, rhs, l2_regularizer)
373  else:
374    return gen_linalg_ops.matrix_solve_ls(
375        matrix, rhs, l2_regularizer, fast=fast, name=name)
376
377
378@tf_export('linalg.eig', 'eig', v1=[])
379@dispatch.add_dispatch_support
380def eig(tensor, name=None):
381  """Computes the eigen decomposition of a batch of matrices.
382
383  The eigenvalues
384  and eigenvectors for a non-Hermitian matrix in general are complex. The
385  eigenvectors are not guaranteed to be linearly independent.
386
387  Computes the eigenvalues and right eigenvectors of the innermost
388  N-by-N matrices in `tensor` such that
389  `tensor[...,:,:] * v[..., :,i] = e[..., i] * v[...,:,i]`, for i=0...N-1.
390
391  Args:
392    tensor: `Tensor` of shape `[..., N, N]`. Only the lower triangular part of
393      each inner inner matrix is referenced.
394    name: string, optional name of the operation.
395
396  Returns:
397    e: Eigenvalues. Shape is `[..., N]`. The eigenvalues are not necessarily
398       ordered.
399    v: Eigenvectors. Shape is `[..., N, N]`. The columns of the inner most
400      matrices contain eigenvectors of the corresponding matrices in `tensor`
401  """
402  if tensor.dtype == dtypes.float32 or tensor.dtype == dtypes.complex64:
403    out_dtype = dtypes.complex64
404  elif tensor.dtype == dtypes.float64 or tensor.dtype == dtypes.complex128:
405    out_dtype = dtypes.complex128
406  e, v = gen_linalg_ops.eig(tensor, Tout=out_dtype, compute_v=True, name=name)
407  return e, v
408
409
410@tf_export('linalg.eigvals', 'eigvals', v1=[])
411@dispatch.add_dispatch_support
412def eigvals(tensor, name=None):
413  """Computes the eigenvalues of one or more matrices.
414
415  Note: If your program backpropagates through this function, you should replace
416  it with a call to tf.linalg.eig (possibly ignoring the second output) to
417  avoid computing the eigen decomposition twice. This is because the
418  eigenvectors are used to compute the gradient w.r.t. the eigenvalues. See
419  _SelfAdjointEigV2Grad in linalg_grad.py.
420
421  Args:
422    tensor: `Tensor` of shape `[..., N, N]`.
423    name: string, optional name of the operation.
424
425  Returns:
426    e: Eigenvalues. Shape is `[..., N]`. The vector `e[..., :]` contains the `N`
427      eigenvalues of `tensor[..., :, :]`.
428  """
429  if tensor.dtype == dtypes.float32 or tensor.dtype == dtypes.complex64:
430    out_dtype = dtypes.complex64
431  elif tensor.dtype == dtypes.float64 or tensor.dtype == dtypes.complex128:
432    out_dtype = dtypes.complex128
433  e, _ = gen_linalg_ops.eig(tensor, Tout=out_dtype, compute_v=False, name=name)
434  return e
435
436
437@tf_export('linalg.eigh', v1=['linalg.eigh', 'self_adjoint_eig'])
438@dispatch.add_dispatch_support
439@deprecation.deprecated_endpoints('self_adjoint_eig')
440def self_adjoint_eig(tensor, name=None):
441  """Computes the eigen decomposition of a batch of self-adjoint matrices.
442
443  Computes the eigenvalues and eigenvectors of the innermost N-by-N matrices
444  in `tensor` such that
445  `tensor[...,:,:] * v[..., :,i] = e[..., i] * v[...,:,i]`, for i=0...N-1.
446
447  Args:
448    tensor: `Tensor` of shape `[..., N, N]`. Only the lower triangular part of
449      each inner inner matrix is referenced.
450    name: string, optional name of the operation.
451
452  Returns:
453    e: Eigenvalues. Shape is `[..., N]`. Sorted in non-decreasing order.
454    v: Eigenvectors. Shape is `[..., N, N]`. The columns of the inner most
455      matrices contain eigenvectors of the corresponding matrices in `tensor`
456  """
457  e, v = gen_linalg_ops.self_adjoint_eig_v2(tensor, compute_v=True, name=name)
458  return e, v
459
460
461@tf_export('linalg.eigvalsh', v1=['linalg.eigvalsh', 'self_adjoint_eigvals'])
462@dispatch.add_dispatch_support
463@deprecation.deprecated_endpoints('self_adjoint_eigvals')
464def self_adjoint_eigvals(tensor, name=None):
465  """Computes the eigenvalues of one or more self-adjoint matrices.
466
467  Note: If your program backpropagates through this function, you should replace
468  it with a call to tf.linalg.eigh (possibly ignoring the second output) to
469  avoid computing the eigen decomposition twice. This is because the
470  eigenvectors are used to compute the gradient w.r.t. the eigenvalues. See
471  _SelfAdjointEigV2Grad in linalg_grad.py.
472
473  Args:
474    tensor: `Tensor` of shape `[..., N, N]`.
475    name: string, optional name of the operation.
476
477  Returns:
478    e: Eigenvalues. Shape is `[..., N]`. The vector `e[..., :]` contains the `N`
479      eigenvalues of `tensor[..., :, :]`.
480  """
481  e, _ = gen_linalg_ops.self_adjoint_eig_v2(tensor, compute_v=False, name=name)
482  return e
483
484
485@tf_export('linalg.svd', v1=['linalg.svd', 'svd'])
486@dispatch.add_dispatch_support
487@deprecation.deprecated_endpoints('svd')
488def svd(tensor, full_matrices=False, compute_uv=True, name=None):
489  r"""Computes the singular value decompositions of one or more matrices.
490
491  Computes the SVD of each inner matrix in `tensor` such that
492  `tensor[..., :, :] = u[..., :, :] * diag(s[..., :, :]) *
493   transpose(conj(v[..., :, :]))`
494
495  ```python
496  # a is a tensor.
497  # s is a tensor of singular values.
498  # u is a tensor of left singular vectors.
499  # v is a tensor of right singular vectors.
500  s, u, v = svd(a)
501  s = svd(a, compute_uv=False)
502  ```
503
504  Args:
505    tensor: `Tensor` of shape `[..., M, N]`. Let `P` be the minimum of `M` and
506      `N`.
507    full_matrices: If true, compute full-sized `u` and `v`. If false
508      (the default), compute only the leading `P` singular vectors.
509      Ignored if `compute_uv` is `False`.
510    compute_uv: If `True` then left and right singular vectors will be
511      computed and returned in `u` and `v`, respectively. Otherwise, only the
512      singular values will be computed, which can be significantly faster.
513    name: string, optional name of the operation.
514
515  Returns:
516    s: Singular values. Shape is `[..., P]`. The values are sorted in reverse
517      order of magnitude, so s[..., 0] is the largest value, s[..., 1] is the
518      second largest, etc.
519    u: Left singular vectors. If `full_matrices` is `False` (default) then
520      shape is `[..., M, P]`; if `full_matrices` is `True` then shape is
521      `[..., M, M]`. Not returned if `compute_uv` is `False`.
522    v: Right singular vectors. If `full_matrices` is `False` (default) then
523      shape is `[..., N, P]`. If `full_matrices` is `True` then shape is
524      `[..., N, N]`. Not returned if `compute_uv` is `False`.
525
526  @compatibility(numpy)
527  Mostly equivalent to numpy.linalg.svd, except that
528    * The order of output  arguments here is `s`, `u`, `v` when `compute_uv` is
529      `True`, as opposed to `u`, `s`, `v` for numpy.linalg.svd.
530    * full_matrices is `False` by default as opposed to `True` for
531       numpy.linalg.svd.
532    * tf.linalg.svd uses the standard definition of the SVD
533      \\(A = U \Sigma V^H\\), such that the left singular vectors of `a` are
534      the columns of `u`, while the right singular vectors of `a` are the
535      columns of `v`. On the other hand, numpy.linalg.svd returns the adjoint
536      \\(V^H\\) as the third output argument.
537  ```python
538  import tensorflow as tf
539  import numpy as np
540  s, u, v = tf.linalg.svd(a)
541  tf_a_approx = tf.matmul(u, tf.matmul(tf.linalg.diag(s), v, adjoint_b=True))
542  u, s, v_adj = np.linalg.svd(a, full_matrices=False)
543  np_a_approx = np.dot(u, np.dot(np.diag(s), v_adj))
544  # tf_a_approx and np_a_approx should be numerically close.
545  ```
546  @end_compatibility
547  """
548  s, u, v = gen_linalg_ops.svd(
549      tensor, compute_uv=compute_uv, full_matrices=full_matrices, name=name)
550  if compute_uv:
551    return math_ops.real(s), u, v
552  else:
553    return math_ops.real(s)
554
555
556# pylint: disable=redefined-builtin
557@tf_export('norm', 'linalg.norm', v1=[])
558@dispatch.add_dispatch_support
559def norm_v2(tensor,
560            ord='euclidean',
561            axis=None,
562            keepdims=None,
563            name=None):
564  r"""Computes the norm of vectors, matrices, and tensors.
565
566  This function can compute several different vector norms (the 1-norm, the
567  Euclidean or 2-norm, the inf-norm, and in general the p-norm for p > 0) and
568  matrix norms (Frobenius, 1-norm, 2-norm and inf-norm).
569
570  Args:
571    tensor: `Tensor` of types `float32`, `float64`, `complex64`, `complex128`
572    ord: Order of the norm. Supported values are `'fro'`, `'euclidean'`,
573      `1`, `2`, `np.inf` and any positive real number yielding the corresponding
574      p-norm. Default is `'euclidean'` which is equivalent to Frobenius norm if
575      `tensor` is a matrix and equivalent to 2-norm for vectors.
576      Some restrictions apply:
577        a) The Frobenius norm `'fro'` is not defined for vectors,
578        b) If axis is a 2-tuple (matrix norm), only `'euclidean'`, '`fro'`, `1`,
579           `2`, `np.inf` are supported.
580      See the description of `axis` on how to compute norms for a batch of
581      vectors or matrices stored in a tensor.
582    axis: If `axis` is `None` (the default), the input is considered a vector
583      and a single vector norm is computed over the entire set of values in the
584      tensor, i.e. `norm(tensor, ord=ord)` is equivalent to
585      `norm(reshape(tensor, [-1]), ord=ord)`.
586      If `axis` is a Python integer, the input is considered a batch of vectors,
587      and `axis` determines the axis in `tensor` over which to compute vector
588      norms.
589      If `axis` is a 2-tuple of Python integers it is considered a batch of
590      matrices and `axis` determines the axes in `tensor` over which to compute
591      a matrix norm.
592      Negative indices are supported. Example: If you are passing a tensor that
593      can be either a matrix or a batch of matrices at runtime, pass
594      `axis=[-2,-1]` instead of `axis=None` to make sure that matrix norms are
595      computed.
596    keepdims: If True, the axis indicated in `axis` are kept with size 1.
597      Otherwise, the dimensions in `axis` are removed from the output shape.
598    name: The name of the op.
599
600  Returns:
601    output: A `Tensor` of the same type as tensor, containing the vector or
602      matrix norms. If `keepdims` is True then the rank of output is equal to
603      the rank of `tensor`. Otherwise, if `axis` is none the output is a scalar,
604      if `axis` is an integer, the rank of `output` is one less than the rank
605      of `tensor`, if `axis` is a 2-tuple the rank of `output` is two less
606      than the rank of `tensor`.
607
608  Raises:
609    ValueError: If `ord` or `axis` is invalid.
610
611  @compatibility(numpy)
612  Mostly equivalent to numpy.linalg.norm.
613  Not supported: ord <= 0, 2-norm for matrices, nuclear norm.
614  Other differences:
615    a) If axis is `None`, treats the flattened `tensor` as a vector
616     regardless of rank.
617    b) Explicitly supports 'euclidean' norm as the default, including for
618     higher order tensors.
619  @end_compatibility
620  """
621  return norm(tensor=tensor,
622              ord=ord,
623              axis=axis,
624              keepdims=keepdims,
625              name=name)
626
627
628# pylint: disable=redefined-builtin
629@tf_export(v1=['norm', 'linalg.norm'])
630@dispatch.add_dispatch_support
631@deprecation.deprecated_args(
632    None, 'keep_dims is deprecated, use keepdims instead', 'keep_dims')
633def norm(tensor,
634         ord='euclidean',
635         axis=None,
636         keepdims=None,
637         name=None,
638         keep_dims=None):
639  r"""Computes the norm of vectors, matrices, and tensors.
640
641  This function can compute several different vector norms (the 1-norm, the
642  Euclidean or 2-norm, the inf-norm, and in general the p-norm for p > 0) and
643  matrix norms (Frobenius, 1-norm, 2-norm and inf-norm).
644
645  Args:
646    tensor: `Tensor` of types `float32`, `float64`, `complex64`, `complex128`
647    ord: Order of the norm. Supported values are 'fro', 'euclidean',
648      `1`, `2`, `np.inf` and any positive real number yielding the corresponding
649      p-norm. Default is 'euclidean' which is equivalent to Frobenius norm if
650      `tensor` is a matrix and equivalent to 2-norm for vectors.
651      Some restrictions apply:
652        a) The Frobenius norm `fro` is not defined for vectors,
653        b) If axis is a 2-tuple (matrix norm), only 'euclidean', 'fro', `1`,
654           `2`, `np.inf` are supported.
655      See the description of `axis` on how to compute norms for a batch of
656      vectors or matrices stored in a tensor.
657    axis: If `axis` is `None` (the default), the input is considered a vector
658      and a single vector norm is computed over the entire set of values in the
659      tensor, i.e. `norm(tensor, ord=ord)` is equivalent to
660      `norm(reshape(tensor, [-1]), ord=ord)`.
661      If `axis` is a Python integer, the input is considered a batch of vectors,
662      and `axis` determines the axis in `tensor` over which to compute vector
663      norms.
664      If `axis` is a 2-tuple of Python integers it is considered a batch of
665      matrices and `axis` determines the axes in `tensor` over which to compute
666      a matrix norm.
667      Negative indices are supported. Example: If you are passing a tensor that
668      can be either a matrix or a batch of matrices at runtime, pass
669      `axis=[-2,-1]` instead of `axis=None` to make sure that matrix norms are
670      computed.
671    keepdims: If True, the axis indicated in `axis` are kept with size 1.
672      Otherwise, the dimensions in `axis` are removed from the output shape.
673    name: The name of the op.
674    keep_dims: Deprecated alias for `keepdims`.
675
676  Returns:
677    output: A `Tensor` of the same type as tensor, containing the vector or
678      matrix norms. If `keepdims` is True then the rank of output is equal to
679      the rank of `tensor`. Otherwise, if `axis` is none the output is a scalar,
680      if `axis` is an integer, the rank of `output` is one less than the rank
681      of `tensor`, if `axis` is a 2-tuple the rank of `output` is two less
682      than the rank of `tensor`.
683
684  Raises:
685    ValueError: If `ord` or `axis` is invalid.
686
687  @compatibility(numpy)
688  Mostly equivalent to numpy.linalg.norm.
689  Not supported: ord <= 0, 2-norm for matrices, nuclear norm.
690  Other differences:
691    a) If axis is `None`, treats the flattened `tensor` as a vector
692     regardless of rank.
693    b) Explicitly supports 'euclidean' norm as the default, including for
694     higher order tensors.
695  @end_compatibility
696  """
697  keepdims = deprecation.deprecated_argument_lookup('keepdims', keepdims,
698                                                    'keep_dims', keep_dims)
699  if keepdims is None:
700    keepdims = False
701
702  is_matrix_norm = ((isinstance(axis, tuple) or isinstance(axis, list)) and
703                    len(axis) == 2)
704  if is_matrix_norm:
705    axis = tuple(axis)
706    if (not isinstance(axis[0], int) or not isinstance(axis[1], int) or
707        axis[0] == axis[1]):
708      raise ValueError(
709          "'axis' must be None, an integer, or a tuple of 2 "
710          f"unique integers, got {axis}")
711    supported_matrix_norms = ['euclidean', 'fro', 1, 2, np.inf]
712    if ord not in supported_matrix_norms:
713      raise ValueError(f"'ord' must be a supported matrix norm in "
714                       f"{supported_matrix_norms}, got {ord}")
715  else:
716    if not (isinstance(axis, int) or axis is None):
717      raise ValueError(
718          "'axis' must be None, an integer, or a "
719          f"tuple of 2 unique integers, got {axis}")
720
721    supported_vector_norms = ['euclidean', 1, 2, np.inf]
722    if (not np.isreal(ord) or ord <= 0) and ord not in supported_vector_norms:
723      raise ValueError(f"'ord' must be a supported vector norm, got {ord}")
724    if axis is not None:
725      axis = (axis,)
726
727  with ops.name_scope(name, 'norm', [tensor]):
728    tensor = ops.convert_to_tensor(tensor)
729
730    if ord in ['fro', 'euclidean', 2, 2.0]:
731      if is_matrix_norm and ord in [2, 2.0]:
732        rank = array_ops.rank(tensor)
733        positive_axis = map_fn.map_fn(
734            lambda i: control_flow_ops.cond(i >= 0, lambda: i, lambda: i + rank
735                                           ), ops.convert_to_tensor(axis))
736        axes = math_ops.range(rank)
737        perm_before = array_ops.concat([
738            gen_array_ops.list_diff(axes, positive_axis, dtypes.int32)[0],
739            positive_axis
740        ],
741                                       axis=0)
742        perm_after = map_fn.map_fn(
743            lambda i: math_ops.cast(
744                array_ops.squeeze(
745                    array_ops.where_v2(math_ops.equal(perm_before, i))),
746                dtype=dtypes.int32), axes)
747        permed = array_ops.transpose(tensor, perm=perm_before)
748        matrix_2_norm = array_ops.expand_dims(
749            math_ops.reduce_max(
750                math_ops.abs(gen_linalg_ops.svd(permed, compute_uv=False)[0]),
751                axis=-1,
752                keepdims=True),
753            axis=-1)
754        result = array_ops.transpose(matrix_2_norm, perm=perm_after)
755      else:
756        result = math_ops.sqrt(
757            math_ops.reduce_sum(
758                tensor * math_ops.conj(tensor), axis, keepdims=True))
759        # TODO(rmlarsen): Replace with the following, once gradients are defined
760        # result = math_ops.reduce_euclidean_norm(tensor, axis, keepdims=True)
761    else:
762      result = math_ops.abs(tensor)
763      if ord == 1:
764        sum_axis = None if axis is None else axis[0]
765        result = math_ops.reduce_sum(result, sum_axis, keepdims=True)
766        if is_matrix_norm:
767          result = math_ops.reduce_max(result, axis[-1], keepdims=True)
768      elif ord == np.inf:
769        if is_matrix_norm:
770          result = math_ops.reduce_sum(result, axis[1], keepdims=True)
771        max_axis = None if axis is None else axis[0]
772        result = math_ops.reduce_max(result, max_axis, keepdims=True)
773      else:
774        # General p-norms (positive p only)
775        result = math_ops.pow(
776            math_ops.reduce_sum(math_ops.pow(result, ord), axis, keepdims=True),
777            1.0 / ord)
778    if not keepdims:
779      result = array_ops.squeeze(result, axis)
780    return result
781
782
783# pylint: enable=invalid-name,redefined-builtin
784