| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ragged/ |
| H A D | ragged_gather_ops.py | 38 batch_dims=0, argument
112 def _gather(params, indices, axis, batch_dims): argument
174 def _batch_gather(params, indices, axis, batch_dims): argument
331 batch_dims=0): argument
341 batch_dims=0, argument
485 batch_dims=0): argument
|
| H A D | ragged_gather_op_test.py | 192 batch_dims=0, argument
396 batch_dims=0): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/math_ops/ |
| H A D | banded_triangular_solve_op_test.py | 27 def _verifySolveAllWays(self, x, y, dtypes, batch_dims=None): argument
40 def _verifySolveAllWaysReal(self, x, y, batch_dims=None): argument
43 def _verifySolveAllWaysComplex(self, x, y, batch_dims=None): argument
51 batch_dims=None, argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/linalg/ |
| H A D | matrix_triangular_solve_op_test.py | 27 def _verifySolveAllWays(self, x, y, dtypes, batch_dims=None): argument
40 def _verifySolveAllWaysReal(self, x, y, batch_dims=None): argument
43 def _verifySolveAllWaysComplex(self, x, y, batch_dims=None): argument
51 batch_dims=None, argument
|
| H A D | matrix_solve_op_test.py | 37 def _verifySolve(self, x, y, batch_dims=None): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | qr_expander.cc | 109 Status House(XlaOp x, XlaOp k, absl::Span<const int64_t> batch_dims, in House()
214 std::vector<int64_t> batch_dims(num_batch_dims); in QrBlock() local
319 PrimitiveType type, absl::Span<const int64_t> batch_dims, XlaOp vs, in CompactWYRepresentation()
394 std::vector<int64_t> batch_dims(num_batch_dims); in BuildQrDecomposition() local
465 std::vector<int64_t> batch_dims(num_batch_dims); in ProductOfElementaryHouseholderReflectors() local
|
| H A D | dot_as_convolution_util.h | 53 std::vector<DimNums> batch_dims; member
|
| H A D | indexed_array_analysis.cc | 981 absl::Span<const int64_t> batch_dims) { in GetOnlyNonContractingNonBatchDim()
1006 absl::Span<const int64_t> batch_dims) { in CanFoldDotIntoIndexedArray()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| H A D | gather_scatter_handler.cc | 170 PartitionedHlo& indices, absl::Span<const int64_t> batch_dims, in PartitionGatherIndexPassthroughPartition()
272 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGatherPassthroughOperand()
376 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGatherTrivialIndexedOperandDimension()
448 std::vector<int64_t> batch_dims; in PartitionGatherTrivialIndexedOperandDimension() local
487 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGatherIndexParallelDimensions()
640 const HloSharding& output_sharding, absl::Span<const int64_t> batch_dims, in PartitionGather()
690 std::vector<int64_t> batch_dims; in HandleGather() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | matmul_utils.cc | 51 const Shape& shape, absl::Span<const int64_t> batch_dims, in GetNonContractingDims()
69 absl::Span<const int64_t> batch_dims, in GetBatchRowColumnShape()
163 const Shape& shape, absl::Span<const int64_t> batch_dims, in For()
184 auto batch_dims = absl::Span<const int64_t>(dims).first(num_batch_dims); in For() local
216 auto batch_dims = (operand_idx == 0) ? dot_dims.lhs_batch_dimensions() in CanFoldTransposeOperandIntoDot() local
|
| H A D | gpu_layout_assignment.cc | 403 absl::Span<const int64_t> batch_dims, absl::Span<const int64_t> row_dims, in SetDotOperandLayout()
425 absl::Span<const int64_t> batch_dims, absl::Span<const int64_t> row_dims, in SetOperandBatchRowsColsLayout()
|
| H A D | cusolver_rewriter.cc | 62 std::vector<int64_t> batch_dims(a_shape.dimensions().begin(), in CreateCholesky() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/array_ops/ |
| H A D | gather_op_test.py | 476 def testBatchDims(self, params, indices, batch_dims, expected=None, argument
594 batch_dims, axis, output_shape): argument
629 def _batchNumpyGather(self, params, indices, axis, batch_dims): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| H A D | svd.cc | 120 std::vector<int64_t> batch_dims(num_batch_dims); in HouseRow() local
185 std::vector<int64_t> batch_dims(num_batch_dims); in HouseCol() local
259 std::vector<int64_t> batch_dims(num_batch_dims); in HouseHolderBidiagonalization() local
461 std::vector<int64_t> batch_dims(num_batch_dims); in OneSidedJacobiUpdate() local
840 std::vector<int64_t> batch_dims(num_batch_dims); in SVD() local
|
| H A D | lu_decomposition.cc | 35 const std::vector<int64_t> batch_dims( in LuDecomposition() local
|
| H A D | self_adjoint_eig.cc | 70 const std::vector<int64_t> batch_dims( in SelfAdjointEig() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/ |
| H A D | gather_spmd_expander.cc | 54 int batch_dims = gather_op.batch_dims(); in ExpandOp() local
174 int batch_dims = gather_op.batch_dims(); in ComputeLayoutForward() local
253 int batch_dims = gather_op.batch_dims(); in ComputeLayoutBackward() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/structured/ |
| H A D | structured_array_ops_test.py | 1113 def testGather(self, params, indices, axis, batch_dims, expected): argument
1150 def testGatherRagged(self, params, indices, axis, batch_dims, expected): argument
1200 indices, axis, batch_dims, argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
| H A D | array_ops.py | 5107 batch_dims=0): # pylint: disable=g-doc-args argument
5317 batch_dims=0, argument
5347 def _batch_gather(params, indices, batch_dims, axis=None): argument
5473 def gather_nd(params, indices, name=None, batch_dims=0): argument
5731 def gather_nd_v2(params, indices, batch_dims=0, name=None): argument
5738 def batch_gather_nd(params, indices, batch_dims, name=None): argument
|
| H A D | array_grad.py | 595 def _GetBatchIndices(params_shape, indices, batch_dims): argument
617 def _BatchGatherGrad(params_shape, values, indices, batch_dims, argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/ops/ |
| H A D | sparse_csr_matrix_ops.cc | 279 ShapeHandle batch_dims; in __anona88575120702() local
432 ShapeHandle batch_dims; in __anona88575120a02() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/cc/gradients/ |
| H A D | array_grad.cc | 632 const Output& indices, int batch_dims) { in GetBatchIndices()
660 Output indices, int batch_dims, Output gather_dim_size) { in BatchGatherGrad()
715 int batch_dims; in GatherV2Grad() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/tests/ |
| H A D | matrix_solve_op_test.py | 54 def testSolve(self, n, nrhs, batch_dims, rhs_batch_dims, adjoint): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/lite/kernels/internal/reference/ |
| H A D | gather.h | 38 int batch_dims = op_params.batch_dims; in Gather() local
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/test/ |
| H A D | legacy_vmap_test.cpp | 94 auto batch_dims = maxBatchDimsAtFront(); in TEST() local
234 auto batch_dims = maxBatchDimsAtFront(); in TEST() local
261 BatchDims batch_dims = { in TEST() local
|