Searched defs:scatter_dimension (Results 1 – 8 of 8) sorted by relevance
1508 const Tensor scatter_dimension = Tensor(0); in TEST() local1564 const Tensor scatter_dimension = Tensor(0); in TEST() local1592 const Tensor scatter_dimension = Tensor(0); in TEST() local
2611 const Tensor* scatter_dimension = c->input_tensor(2); in ReduceScatterShape() local
687 int64_t scatter_dimension() const { return scatter_dimension_; } in scatter_dimension() function
2092 absl::Span<const Shape* const> operand_shapes, int64_t scatter_dimension, in InferReduceScatterShape()
521 int64_t scatter_dimension = proto.dimensions(0); in CreateFromProto() local1294 int64_t scatter_dimension) { in CreateReduceScatter()
901 int64_t scatter_dimension) in HloReduceScatterInstruction()
3018 XlaOp operand, const XlaComputation& computation, int64_t scatter_dimension, in ReduceScatter()4680 int64_t scatter_dimension, int64_t shard_count, in ReduceScatter()
6103 APInt scatter_dimension; in matchAndRewrite() local