Home
last modified time | relevance | path

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

12

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dgather_scatter_handler.cc171 absl::Span<const int64_t> slice_sizes, SpmdPartitioningVisitor* visitor) { in PartitionGatherIndexPassthroughPartition()
273 absl::Span<const int64_t> slice_sizes, PartitionedHlo& operand, in PartitionGatherPassthroughOperand()
377 absl::Span<const int64_t> slice_sizes, PartitionedHlo& operand, in PartitionGatherTrivialIndexedOperandDimension()
488 absl::Span<const int64_t> slice_sizes, PartitionedHlo& operand, in PartitionGatherIndexParallelDimensions()
641 absl::Span<const int64_t> slice_sizes, SpmdPartitioningVisitor* visitor) { in PartitionGather()
777 absl::Span<const int64_t> slice_sizes, SpmdPartitioningVisitor* visitor) { in PartitionScatterPassthroughOperand()
816 absl::Span<const int64_t> slice_sizes, SpmdPartitioningVisitor* visitor) { in PartitionScatterIndexPassthroughPartition()
955 absl::Span<const int64_t> slice_sizes, SpmdPartitioningVisitor* visitor) { in PartitionScatterTrivialIndexedOperandDimension()
1011 absl::Span<const int64_t> slice_sizes, SpmdPartitioningVisitor* visitor) { in PartitionScatter()
1082 std::vector<int64_t> slice_sizes = hlo_sharding_util::GetScatterSliceSize( in HandleScatter() local
/aosp_15_r20/external/executorch/backends/vulkan/runtime/graph/ops/impl/
H A DSDPA.cpp128 std::vector<int64_t> slice_sizes = graph.sizes_of(cache); in get_cache_slice_sizes() local
142 std::vector<int64_t> slice_sizes = get_cache_slice_sizes( in resize_cache_slice_view_node() local
155 std::vector<int64_t> slice_sizes = in add_cache_slice_view_node() local
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A 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 ()()
H A Dsplit_lib_cpu.cc32 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes) { in operator ()()
H A Dsplit_op.cc166 Eigen::DSizes<Eigen::DenseIndex, NDims> slice_sizes; in operator ()() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dslicing.cc216 std::vector<int64_t> slice_sizes(input_shape.rank(), 1); in TorchGather() local
282 std::vector<int64_t> slice_sizes = SpanToVector(input_shape.dimensions()); in TorchIndexSelect() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_instruction.cc583 std::vector<std::vector<int64_t>> slice_sizes; in CreateFromProto() local
799 std::vector<int64_t> slice_sizes(proto.dynamic_slice_sizes_size()); in CreateFromProto() local
844 const auto& slice_sizes = proto.gather_slice_sizes(); in CreateFromProto() local
1336 absl::Span<const std::vector<int64_t>> slice_sizes, in CreateCollectivePermute()
1358 absl::Span<const std::vector<int64_t>> slice_sizes, in CreateCollectivePermuteStart()
1516 absl::Span<const int64_t> slice_sizes) { in CreateDynamicSlice()
1929 absl::Span<const int64_t> slice_sizes, bool indices_are_sorted) { in CreateGather()
4804 int64_t HloInstruction::slice_sizes(int64_t dimension) const { in slice_sizes() function in xla::HloInstruction
H A Dgather_simplifier.cc61 auto slice_sizes = Permute(gather->gather_slice_sizes(), operand_permutation); in ExpandInstruction() local
H A Dhlo_creation_utils.cc163 absl::Span<const int64_t> slice_sizes, const OpMetadata* metadata) { in MakeDynamicSliceHlo()
180 absl::Span<const int64_t> slice_sizes, const OpMetadata* metadata) { in MakeDynamicSliceHlo()
H A Dgather_expander.cc236 absl::Span<const int64_t> slice_sizes, int64_t gather_loop_trip_count, in CreateGatherLoopAccumulatorInitValue()
H A Dhlo_instructions.cc1002 absl::Span<const std::vector<int64_t>> slice_sizes, in HloCollectivePermuteInstruction()
1047 for (const auto& slice_sizes : dynamic_slice_sizes_list()) { in ExtraAttributesToStringImpl() local
2971 absl::Span<const int64_t> slice_sizes) in HloDynamicSliceInstruction()
2981 absl::Span<const int64_t> slice_sizes) in HloDynamicSliceInstruction()
3050 absl::Span<const int64_t> slice_sizes, bool indices_are_sorted) in HloGatherInstruction()
H A Dtriangular_solve_expander.cc71 std::vector<int64_t> slice_sizes(ndims); in DiagonalBlocks() local
H A Dshape_inference.cc2592 absl::Span<const int64_t> slice_sizes, bool allow_scalar_indices) { in InferDynamicSliceShape()
3424 absl::Span<const int64_t> slice_sizes) { in InferGatherShape()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Ddynamic_ops_test.cc119 const std::vector<int64_t>& slice_sizes, in RunR1()
150 const std::vector<int64_t>& slice_sizes, in RunR2()
185 const std::vector<int64_t>& slice_sizes, in RunR3()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/xla/
H A Dget_item_op.cc55 std::vector<int64_t> slice_sizes = {1}; in Compile() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A Dreverse_sequence_op.cc118 auto slice_sizes = input_shape.dim_sizes(); in Compile() local
H A Dgather_scatter_ops.cc38 std::vector<int64_t> slice_sizes; in Compile() local
H A Dfft_ops.cc62 absl::InlinedVector<int64_t, 4> slice_sizes = input_shape.dim_sizes(); in Compile() local
H A Dgather_op.cc124 std::vector<int64_t> slice_sizes; in XlaGather() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
H A Dmlir_hlo_builder.cc384 absl::Span<const int64_t> slice_sizes, bool indices_are_sorted) { in GatherInternal()
664 absl::Span<const int64_t> slice_sizes) { in DynamicSliceInternal()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/
H A Dhlo_function_importer.cc698 std::vector<int64_t> slice_sizes( in ImportInstructionImpl() local
712 std::vector<int64_t> slice_sizes( in ImportInstructionImpl() local
/aosp_15_r20/external/mesa3d/src/gallium/drivers/d3d12/
H A Dd3d12_video_enc_h264.cpp418 std::vector<int> slice_sizes(picture->num_slice_descriptors); in d3d12_video_encoder_negotiate_current_h264_slices_configuration() local
H A Dd3d12_video_enc_hevc.cpp519 std::vector<int> slice_sizes(picture->num_slice_descriptors); in d3d12_video_encoder_negotiate_current_hevc_slices_configuration() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dcollection_ops_util.cc303 auto slice_sizes = GetR1Const(per_slice_shape, builder, loc); in ScatterAccumulateElements() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
H A Dlegalize_tf.cc266 auto slice_sizes = llvm::to_vector<4>(type.getShape()); in DynamicSliceInMinorDims() local
869 DenseIntElementsAttr slice_sizes) { in CanBeTranslatedToDynamicSlice()
896 Value input, Value start_indices, DenseIntElementsAttr slice_sizes, in TFSliceSizes2HLOSliceSizes()
1501 SmallVector<int64_t, 4> slice_sizes; in matchAndRewrite() local
2034 Shape slice_sizes(input_shape.begin(), input_shape.end()); in matchAndRewrite() local
3736 SmallVector<int64_t, 4> slice_sizes; in rewriteWithUnknownBegin() local
5958 SmallVector<int64_t, 4> slice_sizes(input_shape.begin(), input_shape.end()); in matchAndRewrite() local

12