Home
last modified time | relevance | path

Searched defs:batch_dims (Results 1 – 25 of 56) sorted by relevance

123

/external/tensorflow/tensorflow/python/ops/ragged/
Dragged_gather_ops.py38 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
Dragged_gather_op_test.py192 batch_dims=0, argument
396 batch_dims=0): argument
/external/tensorflow/tensorflow/python/kernel_tests/math_ops/
Dbanded_triangular_solve_op_test.py27 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
/external/tensorflow/tensorflow/python/kernel_tests/linalg/
Dmatrix_triangular_solve_op_test.py27 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
Dmatrix_solve_op_test.py37 def _verifySolve(self, x, y, batch_dims=None): argument
Dmatrix_exponential_op_test.py218 def _TestRandomSmall(dtype, batch_dims, size): argument
/external/tensorflow/tensorflow/python/kernel_tests/array_ops/
Dgather_op_test.py476 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
/external/tensorflow/tensorflow/compiler/xla/service/
Dqr_expander.cc109 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
Ddot_as_convolution_util.h53 std::vector<DimNums> batch_dims; member
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dmatmul_utils.cc51 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
Dgpu_layout_assignment.cc403 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()
Dcusolver_rewriter.cc62 std::vector<int64_t> batch_dims(a_shape.dimensions().begin(), in CreateCholesky() local
Dmatmul_utils_test.cc41 std::vector<int64_t> batch_dims; member
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dlu_decomposition.cc35 const std::vector<int64_t> batch_dims( in LuDecomposition() local
Dself_adjoint_eig.cc70 const std::vector<int64_t> batch_dims( in SelfAdjointEig() local
Dsvd.cc120 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
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dgather_scatter_handler.cc170 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
/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dgather_spmd_expander.cc54 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
/external/tensorflow/tensorflow/compiler/tests/
Dmatrix_solve_op_test.py54 def testSolve(self, n, nrhs, batch_dims, rhs_batch_dims, adjoint): argument
/external/tensorflow/tensorflow/lite/kernels/internal/reference/
Dgather.h38 int batch_dims = op_params.batch_dims; in Gather() local
/external/tensorflow/tensorflow/python/ops/structured/
Dstructured_array_ops_test.py1113 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
/external/tensorflow/tensorflow/core/ops/
Dsparse_csr_matrix_ops.cc279 ShapeHandle batch_dims; in __anon6adce4610702() local
432 ShapeHandle batch_dims; in __anon6adce4610a02() local
/external/tensorflow/tensorflow/core/kernels/
Dgather_op.cc92 int32_t batch_dims = batch_dims_; in Compute() local
Dgather_op_test.cc43 void MakeOp(DataType data_type, DataType index_type, int batch_dims = 0) { in MakeOp()
/external/tensorflow/tensorflow/lite/kernels/
Dgather.cc90 int batch_dims = params->batch_dims; in Prepare() local

123