Searched defs:offset_dims (Results 1 – 11 of 11) sorted by relevance
67 std::vector<int64_t> offset_dims(operand_rank); in ExpandInstruction() local
255 HloInstruction* accumulator, absl::Span<const int64_t> offset_dims, in PermuteBatchAndOffsetDims()
1070 std::vector<int64_t> offset_dims(dnums.offset_dims().begin(), in GatherOutputShardingFromDataOperand() local1084 std::vector<int64_t> offset_dims(dnums.offset_dims().begin(), in GatherDataOperandShardingFromOutput() local
3062 std::string offset_dims = in GatherDimensionNumbersToString() local3080 absl::Span<const int64_t> offset_dims, in MakeGatherDimNumbers()
2804 optional<std::vector<int64_t>> offset_dims; in CreateInstruction() local
48 std::vector<int64_t> offset_dims(dnums.offset_dims().begin(), in ConvertGatherDimensionNumbers() local
413 auto offset_dims = input.getOffsetDims(); in Convert_dimension_numbers() local
535 repeated int64 offset_dims = 1; field
2939 auto offset_dims = gather_op.dimension_numbers().getOffsetDims(); in matchAndRewrite() local
1541 SmallVector<int64_t, 4> offset_dims; in matchAndRewrite() local
4838 std::vector<int64_t> offset_dims = {x_dims[is_nhwc ? 3 : 1]}; in TEST_F() local