| /external/tensorflow/tensorflow/core/kernels/linalg/ |
| D | tridiagonal_matmul_op_gpu.cu.cc | 39 const Scalar* __restrict__ maindiag, in TridiagonalMatMulKernel() argument
45 Scalar result = maindiag[row_id] * rhs[i]; in TridiagonalMatMulKernel()
64 const Tensor& maindiag = context->input(1); in Compute() local
85 maindiag.flat<Scalar>().data(), subdiag.flat<Scalar>().data(), in Compute()
|
| D | tridiagonal_matmul_op.cc | 90 const auto& maindiag = inputs[1].row(0); in ComputeMatrix() local
111 result.array() = rhs.array().colwise() * maindiag.transpose().array(); in ComputeMatrix()
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | tridiagonal_matmul_op_test.py | 43 maindiag, argument
50 diags_compact = np.stack([superdiag_extended, maindiag, subdiag_extended])
51 diags_matrix = np.diag(superdiag, 1) + np.diag(maindiag, 0) + np.diag(
55 constant_op.constant(maindiag, dtype),
94 def _makeTridiagonalMatrix(self, superdiag, maindiag, subdiag): argument
99 main_part = array_ops.matrix_diag(maindiag)
146 maindiag = self._randomComplexArray((b, m))
150 np.diag(maindiag[i], 0) + \
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | linalg_ops.cc | 286 ShapeHandle maindiag; in TridiagonalMatMulShapeFn() local
292 TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 2, &maindiag)); in TridiagonalMatMulShapeFn()
302 TF_RETURN_IF_ERROR(c->Subshape(maindiag, 0, -2, &maindiag_batch_shape)); in TridiagonalMatMulShapeFn()
305 TF_RETURN_IF_ERROR(c->Merge(superdiag, maindiag, &superdiag)); in TridiagonalMatMulShapeFn()
312 TF_RETURN_IF_ERROR(c->Merge(superdiag, maindiag, &maindiag)); in TridiagonalMatMulShapeFn()
313 TF_RETURN_IF_ERROR(c->Merge(subdiag, maindiag, &maindiag)); in TridiagonalMatMulShapeFn()
317 DimensionHandle m_lhs = c->Dim(maindiag, -1); in TridiagonalMatMulShapeFn()
323 TF_RETURN_IF_ERROR(c->WithValue(c->Dim(maindiag, -2), 1, &unused)); in TridiagonalMatMulShapeFn()
|
| D | ops.pbtxt | 56781 name: "maindiag"
|
| /external/tensorflow/tensorflow/python/ops/linalg/ |
| D | linalg_impl.py | 555 superdiag, maindiag, subdiag = diagonals
556 if (not subdiag.shape[:-1].is_compatible_with(maindiag.shape[:-1]) or
557 not superdiag.shape[:-1].is_compatible_with(maindiag.shape[:-1])):
561 subdiag.shape, maindiag.shape, superdiag.shape))
563 m = tensor_shape.dimension_value(maindiag.shape[-1])
579 diagonals = array_ops.stack((superdiag, maindiag, subdiag), axis=-2)
701 maindiag = diagonals[..., 1, :]
704 superdiag, maindiag, subdiag = diagonals
715 maindiag = diags[..., 1, :]
723 maindiag = array_ops.expand_dims(maindiag, -2)
[all …]
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | TridiagonalMatMul.pbtxt | 8 name: "maindiag"
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/ |
| D | TridiagonalMatMul.pbtxt | 8 name: "maindiag"
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_TridiagonalMatMul.pbtxt | 12 name: "maindiag"
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.raw_ops.pbtxt | 5157 …argspec: "args=[\'superdiag\', \'maindiag\', \'subdiag\', \'rhs\', \'name\'], varargs=None, keywor…
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.raw_ops.pbtxt | 5157 …argspec: "args=[\'superdiag\', \'maindiag\', \'subdiag\', \'rhs\', \'name\'], varargs=None, keywor…
|
| /external/tensorflow/tensorflow/go/op/ |
| D | wrappers.go | 15036 func TridiagonalMatMul(scope *Scope, superdiag tf.Output, maindiag tf.Output, subdiag tf.Output, rh…
15043 superdiag, maindiag, subdiag, rhs,
|