Searched refs:output_batch_shape (Results 1 – 12 of 12) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | matrix_triangular_solve_op.cc | 89 TensorShape lhs_broadcast_shape(broadcast_helper.output_batch_shape()); in Broadcast()
98 TensorShape rhs_broadcast_shape(broadcast_helper.output_batch_shape()); in Broadcast()
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | linalg_ops.cc | 74 ShapeHandle output_batch_shape; in BandedTriangularSolveShapeFn() local
79 c, lhs_batch_shape, rhs_batch_shape, true, &output_batch_shape)); in BandedTriangularSolveShapeFn()
87 c->Concatenate(output_batch_shape, c->Matrix(m, c->Dim(rhs, -1)), &out)); in BandedTriangularSolveShapeFn()
140 ShapeHandle output_batch_shape; in MatrixTriangularSolveShapeFn() local
145 c, lhs_batch_shape, rhs_batch_shape, true, &output_batch_shape)); in MatrixTriangularSolveShapeFn()
153 c->Concatenate(output_batch_shape, c->Matrix(m, c->Dim(rhs, -1)), &out)); in MatrixTriangularSolveShapeFn()
|
| /external/tensorflow/tensorflow/core/util/ |
| D | matmul_bcast.h | 72 const TensorShape& output_batch_shape() const { return output_batch_shape_; } in output_batch_shape() function
|
| D | matmul_bcast_test.cc | 31 &ret, "[", absl::StrJoin(b.output_batch_shape().dim_sizes(), ","), "]"); in MatMulBCastToStr()
|
| /external/tensorflow/tensorflow/core/kernels/mkl/ |
| D | mkl_batch_matmul_op.cc | 104 TensorShape out_shape = bcast.output_batch_shape(); in Compute()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | unroll_batch_matmul.cc | 297 const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes(); in matchAndRewrite()
|
| /external/tensorflow/tensorflow/lite/toco/graph_transformations/ |
| D | unroll_batch_matmul.cc | 244 const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes(); in Run()
|
| /external/tensorflow/tensorflow/core/kernels/linalg/ |
| D | banded_triangular_solve_op.cc | 228 TensorShape out_shape = bcast.output_batch_shape(); in Compute()
|
| D | matrix_triangular_solve_op_impl.h | 173 TensorShape out_shape = bcast.output_batch_shape();
|
| D | einsum_op_impl.h | 564 TensorShape output_shape = bcast.output_batch_shape(); in ContractOperands()
|
| /external/tensorflow/tensorflow/core/framework/ |
| D | common_shape_fns.cc | 331 ShapeHandle output_batch_shape; in BatchMatMulV2Shape() local
336 c, a_batch_shape, b_batch_shape, true, &output_batch_shape)); in BatchMatMulV2Shape()
340 output_batch_shape, c->Matrix(output_rows, output_cols), &output_shape)); in BatchMatMulV2Shape()
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | matmul_op_impl.h | 682 TensorShape out_shape = bcast.output_batch_shape();
|