Home
last modified time | relevance | path

Searched refs:a_batch_dims (Results 1 – 4 of 4) sorted by relevance

/external/tensorflow/tensorflow/core/ops/
Dsparse_csr_matrix_ops.cc277 ShapeHandle a_batch_dims; in __anon59a9d6c30702() local
280 TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_batch_dims)); in __anon59a9d6c30702()
282 TF_RETURN_IF_ERROR(c->Merge(a_batch_dims, b_batch_dims, &batch_dims)); in __anon59a9d6c30702()
430 ShapeHandle a_batch_dims; in __anon59a9d6c30a02() local
433 TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_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/
Dcommon_shape_fns.cc361 ShapeHandle a_batch_dims; in BatchMatMulShape() local
364 TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_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/
Dshape_inference.cc1994 absl::Span<const int64> a_batch_dims(a.dimensions()); in InferTriangularSolveShape() local
1996 a_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/
Dhlo_ops.cc2771 auto a_batch_dims = a_type.getShape().drop_back(2); in Verify() local
2773 if (a_batch_dims != b_batch_dims) in Verify()