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_t split_dimension_; member in tensorflow::__anonbe76eefc0111::AllToAllOp
949 split_dimension_(split_dimension) {} in HloAllToAllInstruction()962 if (split_dimension_) { in ToProto()963 proto.add_dimensions(*split_dimension_); in ToProto()972 if (split_dimension_) { in ExtraAttributesToStringImpl()973 result.push_back(StrCat("dimensions={", *split_dimension_, "}")); in ExtraAttributesToStringImpl()985 split_dimension_ == casted_other.split_dimension(); in IdenticalSlowPathIgnoringChannelIdValues()
730 std::optional<int64_t> split_dimension() const { return split_dimension_; } in split_dimension()731 void set_split_dimension(int64_t dim) { split_dimension_ = dim; } in set_split_dimension()753 std::optional<int64_t> split_dimension_; variable