Home
last modified time | relevance | path

Searched defs:slice_sizes (Results 1 – 25 of 35) sorted by relevance

12

/external/tensorflow/tensorflow/core/kernels/
Dsplit_lib_cpu.cc32 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes) { in operator ()()
Dsplit_lib_gpu.cu.cc37 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes) { in operator ()()
46 const Eigen::DSizes<Eigen::DenseIndex, 2>& slice_sizes) { in operator ()()
Dsplit_op.cc166 Eigen::DSizes<Eigen::DenseIndex, NDims> slice_sizes; in operator ()() local
/external/tensorflow/tensorflow/core/tpu/kernels/xla/
Dget_item_op.cc55 std::vector<int64> slice_sizes = {1}; in Compile() local
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dreverse_sequence_op.cc117 auto slice_sizes = input_shape.dim_sizes(); in Compile() local
Dgather_scatter_ops.cc38 std::vector<int64> slice_sizes; in Compile() local
Dfft_ops.cc62 absl::InlinedVector<int64, 4> slice_sizes = input_shape.dim_sizes(); in Compile() local
Dgather_op.cc123 std::vector<int64> slice_sizes; in XlaGather() local
Dstrided_slice_op.cc106 std::vector<int64> slice_sizes(input_shape.dims(), 1); in EmitDynamicSlice() local
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dslicing.cc215 std::vector<int64> slice_sizes(input_shape.rank(), 1); in TorchGather() local
279 std::vector<int64> slice_sizes = SpanToVector(input_shape.dimensions()); in TorchIndexSelect() local
Dmatrix.cc151 std::vector<int64> slice_sizes; in GetMatrixDiagonalViaGather() local
/external/tensorflow/tensorflow/compiler/xla/service/
Dgather_expander.cc235 absl::Span<const int64> slice_sizes, int64_t gather_loop_trip_count, in CreateGatherLoopAccumulatorInitValue()
Dhlo_creation_utils.cc137 absl::Span<const int64> slice_sizes) { in MakeDynamicSliceHlo()
152 absl::Span<const int64> slice_sizes) { in MakeDynamicSliceHlo()
Dhlo_instruction.cc527 std::vector<std::vector<int64_t>> slice_sizes; in CreateFromProto() local
741 std::vector<int64> slice_sizes(proto.dynamic_slice_sizes_size()); in CreateFromProto() local
1239 absl::Span<const std::vector<int64_t>> slice_sizes, in CreateCollectivePermute()
1261 absl::Span<const std::vector<int64_t>> slice_sizes, in CreateCollectivePermuteStart()
1418 absl::Span<const int64> slice_sizes) { in CreateDynamicSlice()
1805 absl::Span<const int64> slice_sizes, bool indices_are_sorted) { in CreateGather()
4526 int64 HloInstruction::slice_sizes(int64_t dimension) const { in slice_sizes() function in xla::HloInstruction
Dtriangular_solve_expander.cc76 std::vector<int64> slice_sizes(ndims); in DiagonalBlocks() local
Dhlo_instructions.cc838 absl::Span<const std::vector<int64_t>> slice_sizes, in HloCollectivePermuteInstruction()
879 for (const auto& slice_sizes : dynamic_slice_sizes_list()) { in ExtraAttributesToStringImpl() local
2812 absl::Span<const int64> slice_sizes) in HloDynamicSliceInstruction()
2822 absl::Span<const int64> slice_sizes) in HloDynamicSliceInstruction()
2890 absl::Span<const int64> slice_sizes, bool indices_are_sorted) in HloGatherInstruction()
Dshape_inference.cc2589 absl::Span<const int64> slice_sizes, bool allow_scalar_indices) { in InferDynamicSliceShape()
3442 absl::Span<const int64> slice_sizes) { in InferGatherShape()
Dindexed_array_analysis.cc255 absl::Span<const int64> slice_sizes, Array* source, Array* indices) { in ComputeArrayForGather()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dhlo_function_importer.cc375 std::vector<int64_t> slice_sizes( in ImportInstructionImpl() local
387 std::vector<int64_t> slice_sizes( in ImportInstructionImpl() local
/external/tensorflow/tensorflow/compiler/xla/tests/
Ddynamic_ops_test.cc120 const std::vector<int64>& slice_sizes, in RunR1()
151 const std::vector<int64>& slice_sizes, in RunR2()
186 const std::vector<int64>& slice_sizes, in RunR3()
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.cc307 absl::Span<const int64> slice_sizes, bool indices_are_sorted) { in GatherInternal()
536 absl::Span<const int64> slice_sizes) { in DynamicSliceInternal()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dcollection_ops_util.cc303 auto slice_sizes = GetR1Const(per_slice_shape, builder, loc); in ScatterAccumulateElements() local
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/
Dlhlo_legalize_to_affine.cc293 SmallVector<int64_t, 4> slice_sizes; in matchAndRewrite() local
/external/tensorflow/tensorflow/compiler/tf2xla/python/
Dxla.py550 def gather(operand, start_indices, dimension_numbers, slice_sizes, argument
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf.cc312 auto slice_sizes = llvm::to_vector<4>(type.getShape()); in DynamicSliceInMinorDims() local
943 DenseIntElementsAttr slice_sizes) { in CanBeTranslatedToDynamicSlice()
970 Value input, Value start_indices, DenseIntElementsAttr slice_sizes, in TFSliceSizes2HLOSliceSizes()
1165 SmallVector<int64_t, 4> slice_sizes; in matchAndRewrite() local
1677 SmallVector<int64_t, 4> slice_sizes; in matchAndRewrite() local
2241 Shape slice_sizes(input_shape.begin(), input_shape.end()); in matchAndRewrite() local
3935 SmallVector<int64_t, 4> slice_sizes(slicing_dim_size, 1); in rewriteWithUnknownBegin() local
6241 SmallVector<int64_t, 4> slice_sizes(input_shape.begin(), input_shape.end()); in matchAndRewrite() local

12