Home
last modified time | relevance | path

Searched refs:output_batch_shape (Results 1 – 11 of 11) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dmatrix_triangular_solve_op.cc81 TensorShape lhs_broadcast_shape(broadcast_helper.output_batch_shape()); in Broadcast()
90 TensorShape rhs_broadcast_shape(broadcast_helper.output_batch_shape()); in Broadcast()
/external/tensorflow/tensorflow/core/util/
Dmatmul_bcast.h42 const TensorShape& output_batch_shape() const { return output_shape_; } in output_batch_shape() function
Dmatmul_bcast_test.cc31 &ret, "[", absl::StrJoin(b.output_batch_shape().dim_sizes(), ","), "]"); in MatMulBCastToStr()
/external/tensorflow/tensorflow/core/ops/
Dlinalg_ops.cc97 ShapeHandle output_batch_shape; in MatrixTriangularSolveShapeFn() local
102 c, lhs_batch_shape, rhs_batch_shape, true, &output_batch_shape)); in MatrixTriangularSolveShapeFn()
110 c->Concatenate(output_batch_shape, c->Matrix(m, c->Dim(rhs, -1)), &out)); in MatrixTriangularSolveShapeFn()
/external/tensorflow/tensorflow/lite/toco/graph_transformations/
Dunroll_batch_matmul.cc243 const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes(); in Run()
/external/tensorflow/tensorflow/compiler/mlir/lite/transforms/
Dunroll_batch_matmul.cc302 const auto& result_batch_shape = bcast.output_batch_shape().dim_sizes(); in matchAndRewrite()
/external/tensorflow/tensorflow/core/kernels/
Dmkl_batch_matmul_op.cc106 TensorShape out_shape = bcast.output_batch_shape(); in Compute()
Dmatrix_triangular_solve_op_impl.h173 TensorShape out_shape = bcast.output_batch_shape();
Deinsum_op_impl.h581 TensorShape output_shape = bcast.output_batch_shape(); in ContractOperands()
Dbatch_matmul_op_impl.h609 TensorShape out_shape = bcast.output_batch_shape();
/external/tensorflow/tensorflow/core/framework/
Dcommon_shape_fns.cc440 ShapeHandle output_batch_shape; in BatchMatMulV2Shape() local
445 c, a_batch_shape, b_batch_shape, true, &output_batch_shape)); in BatchMatMulV2Shape()
449 output_batch_shape, c->Matrix(output_rows, output_cols), &output_shape)); in BatchMatMulV2Shape()