Searched refs:split_dimension_ (Results 1 – 3 of 3) sorted by relevance
71 OP_REQUIRES_OK(ctx, ctx->GetAttr("split_dimension", &split_dimension_)); in AllToAllOp()84 0, xla::AllToAll(ctx->Input(0), split_dimension_, concat_dimension_, in Compile()89 int64 split_dimension_; member in tensorflow::__anon2b5c5e460111::AllToAllOp
729 split_dimension_(split_dimension) {} in HloAllToAllInstruction()742 if (split_dimension_) { in ToProto()743 proto.add_dimensions(*split_dimension_); in ToProto()752 if (split_dimension_) { in ExtraAttributesToStringImpl()753 result.push_back(StrCat("dimensions={", *split_dimension_, "}")); in ExtraAttributesToStringImpl()765 split_dimension_ == casted_other.split_dimension(); in IdenticalSlowPathIgnoringChannelIdValues()
478 absl::optional<int64> split_dimension() const { return split_dimension_; } in split_dimension()479 void set_split_dimension(int64 dim) { split_dimension_ = dim; } in set_split_dimension()497 absl::optional<int64> split_dimension_; variable