Searched refs:b_batch_dims (Results 1 – 4 of 4) sorted by relevance
/external/tensorflow/tensorflow/core/ops/ |
D | sparse_csr_matrix_ops.cc | 278 ShapeHandle b_batch_dims; in __anon59a9d6c30702() local 281 TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_dims)); in __anon59a9d6c30702() 282 TF_RETURN_IF_ERROR(c->Merge(a_batch_dims, b_batch_dims, &batch_dims)); in __anon59a9d6c30702() 431 ShapeHandle b_batch_dims; in __anon59a9d6c30a02() local 434 TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_dims)); in __anon59a9d6c30a02() 435 TF_RETURN_IF_ERROR(c->Merge(a_batch_dims, b_batch_dims, &batch_dims)); in __anon59a9d6c30a02()
|
/external/tensorflow/tensorflow/core/framework/ |
D | common_shape_fns.cc | 362 ShapeHandle b_batch_dims; in BatchMatMulShape() local 365 TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_dims)); in BatchMatMulShape() 366 TF_RETURN_IF_ERROR(c->Merge(a_batch_dims, b_batch_dims, &batch_dims)); in BatchMatMulShape()
|
/external/tensorflow/tensorflow/compiler/xla/service/ |
D | shape_inference.cc | 1995 absl::Span<const int64> b_batch_dims(b.dimensions()); in InferTriangularSolveShape() local 1997 b_batch_dims.remove_suffix(2); in InferTriangularSolveShape() 1998 if (a_batch_dims != b_batch_dims) { in InferTriangularSolveShape()
|
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/IR/ |
D | hlo_ops.cc | 2772 auto b_batch_dims = b_type.getShape().drop_back(2); in Verify() local 2773 if (a_batch_dims != b_batch_dims) in Verify()
|