V2 ops always align the diagonals to the left (LEFT_LEFT) in the compact format. V3 ops support 4 alignments: RIGHT_LEFT, LEFT_RIGHT, LEFT_LEFT, and RIGHT_RIGHT. We would like to use RIGHT_LEFT as the default alignment. This contradicts with v2's behavior so we need new a version. V2 has never been exposed to the public APIs. We will skip V2 and go from V1 to V3 directly. V3 features are currently under forward compatibility guards and will be enabled automatically in ~3 weeks from now. This commit contains - V3 API definitions. - Modifications to C++ Matrix{Diag,SetDiag,DiagPart}Op kernels (CPU, GPU, XLA) and shape inference functions to support v3. - Additional tests and gradient implementations in Python for v3. - Pfor and TFLite TOCO converters for v3. - The TFLite MLIR converter for MatrixDiagV3 is intentionally left out because of an MLIR test infrastructure issue and will be added in a separate commit. Notes: - Python changes cannot be in a separate follow-up commit because all kernel tests are in Python. (No C++ tests.) - All three ops have to be in the same commit because their gradients call each other. PiperOrigin-RevId: 280527550 Change-Id: I88e91abab5c4b50419204807ede4fa60657f048a
260 lines
9.1 KiB
Plaintext
260 lines
9.1 KiB
Plaintext
path: "tensorflow.linalg"
|
|
tf_module {
|
|
member {
|
|
name: "LinearOperator"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorAdjoint"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorBlockDiag"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorCirculant"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorCirculant2D"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorCirculant3D"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorComposition"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorDiag"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorFullMatrix"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorHouseholder"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorIdentity"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorInversion"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorKronecker"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorLowRankUpdate"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorLowerTriangular"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorPermutation"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorScaledIdentity"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorToeplitz"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member {
|
|
name: "LinearOperatorZeros"
|
|
mtype: "<type \'type\'>"
|
|
}
|
|
member_method {
|
|
name: "adjoint"
|
|
argspec: "args=[\'matrix\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "band_part"
|
|
argspec: "args=[\'input\', \'num_lower\', \'num_upper\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "cholesky"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "cholesky_solve"
|
|
argspec: "args=[\'chol\', \'rhs\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "cross"
|
|
argspec: "args=[\'a\', \'b\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "det"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "diag"
|
|
argspec: "args=[\'diagonal\', \'name\', \'k\', \'num_rows\', \'num_cols\', \'padding_value\', \'align\'], varargs=None, keywords=None, defaults=[\'diag\', \'0\', \'-1\', \'-1\', \'0\', \'RIGHT_LEFT\'], "
|
|
}
|
|
member_method {
|
|
name: "diag_part"
|
|
argspec: "args=[\'input\', \'name\', \'k\', \'padding_value\', \'align\'], varargs=None, keywords=None, defaults=[\'diag_part\', \'0\', \'0\', \'RIGHT_LEFT\'], "
|
|
}
|
|
member_method {
|
|
name: "eig"
|
|
argspec: "args=[\'tensor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "eigh"
|
|
argspec: "args=[\'tensor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "eigvals"
|
|
argspec: "args=[\'tensor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "eigvalsh"
|
|
argspec: "args=[\'tensor\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "einsum"
|
|
argspec: "args=[\'equation\'], varargs=inputs, keywords=kwargs, defaults=None"
|
|
}
|
|
member_method {
|
|
name: "expm"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "eye"
|
|
argspec: "args=[\'num_rows\', \'num_columns\', \'batch_shape\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"<dtype: \'float32\'>\", \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "global_norm"
|
|
argspec: "args=[\'t_list\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "inv"
|
|
argspec: "args=[\'input\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "l2_normalize"
|
|
argspec: "args=[\'x\', \'axis\', \'epsilon\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'1e-12\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "logdet"
|
|
argspec: "args=[\'matrix\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "logm"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "lstsq"
|
|
argspec: "args=[\'matrix\', \'rhs\', \'l2_regularizer\', \'fast\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'True\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "lu"
|
|
argspec: "args=[\'input\', \'output_idx_type\', \'name\'], varargs=None, keywords=None, defaults=[\"<dtype: \'int32\'>\", \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "lu_matrix_inverse"
|
|
argspec: "args=[\'lower_upper\', \'perm\', \'validate_args\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "lu_reconstruct"
|
|
argspec: "args=[\'lower_upper\', \'perm\', \'validate_args\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "lu_solve"
|
|
argspec: "args=[\'lower_upper\', \'perm\', \'rhs\', \'validate_args\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "matmul"
|
|
argspec: "args=[\'a\', \'b\', \'transpose_a\', \'transpose_b\', \'adjoint_a\', \'adjoint_b\', \'a_is_sparse\', \'b_is_sparse\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'False\', \'False\', \'False\', \'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "matrix_rank"
|
|
argspec: "args=[\'a\', \'tol\', \'validate_args\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "matrix_transpose"
|
|
argspec: "args=[\'a\', \'name\', \'conjugate\'], varargs=None, keywords=None, defaults=[\'matrix_transpose\', \'False\'], "
|
|
}
|
|
member_method {
|
|
name: "matvec"
|
|
argspec: "args=[\'a\', \'b\', \'transpose_a\', \'adjoint_a\', \'a_is_sparse\', \'b_is_sparse\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'False\', \'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "norm"
|
|
argspec: "args=[\'tensor\', \'ord\', \'axis\', \'keepdims\', \'name\'], varargs=None, keywords=None, defaults=[\'euclidean\', \'None\', \'None\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "normalize"
|
|
argspec: "args=[\'tensor\', \'ord\', \'axis\', \'name\'], varargs=None, keywords=None, defaults=[\'euclidean\', \'None\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "pinv"
|
|
argspec: "args=[\'a\', \'rcond\', \'validate_args\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "qr"
|
|
argspec: "args=[\'input\', \'full_matrices\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "set_diag"
|
|
argspec: "args=[\'input\', \'diagonal\', \'name\', \'k\', \'align\'], varargs=None, keywords=None, defaults=[\'set_diag\', \'0\', \'RIGHT_LEFT\'], "
|
|
}
|
|
member_method {
|
|
name: "slogdet"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "solve"
|
|
argspec: "args=[\'matrix\', \'rhs\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "sqrtm"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "svd"
|
|
argspec: "args=[\'tensor\', \'full_matrices\', \'compute_uv\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'True\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "tensor_diag"
|
|
argspec: "args=[\'diagonal\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "tensor_diag_part"
|
|
argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "tensordot"
|
|
argspec: "args=[\'a\', \'b\', \'axes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "trace"
|
|
argspec: "args=[\'x\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
|
|
}
|
|
member_method {
|
|
name: "triangular_solve"
|
|
argspec: "args=[\'matrix\', \'rhs\', \'lower\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'False\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "tridiagonal_matmul"
|
|
argspec: "args=[\'diagonals\', \'rhs\', \'diagonals_format\', \'name\'], varargs=None, keywords=None, defaults=[\'compact\', \'None\'], "
|
|
}
|
|
member_method {
|
|
name: "tridiagonal_solve"
|
|
argspec: "args=[\'diagonals\', \'rhs\', \'diagonals_format\', \'transpose_rhs\', \'conjugate_rhs\', \'name\', \'partial_pivoting\'], varargs=None, keywords=None, defaults=[\'compact\', \'False\', \'False\', \'None\', \'True\'], "
|
|
}
|
|
}
|