Searched defs:split_dimension (Results 1 – 13 of 13) sorted by relevance
| /external/tensorflow/tensorflow/core/ops/ |
| D | tpu_cross_replica_ops.cc | 42 int split_dimension; in __anonc85edbdb0102() local
|
| D | array_ops.cc | 574 DimensionHandle split_dimension; in __anon42d741190a02() local
610 DimensionHandle split_dimension; in __anon42d741190b02() local
|
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 129 def split(cls, tensor, split_dimension, num_devices, input_shape=None): argument
299 split_dimension, argument
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 47 split_dimension, argument
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 50 const int split_dimension, in CreateSplitOp()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_instructions.h | 478 absl::optional<int64> split_dimension() const { return split_dimension_; } in split_dimension() function
|
| D | hlo_instruction.cc | 458 absl::optional<int64> split_dimension; in CreateFromProto() local
1066 const absl::optional<int64>& split_dimension) { in CreateAllToAll()
|
| D | shape_inference.cc | 2062 const Shape& shape, int64 split_dimension, int64 concat_dimension, in InferAllToAllShape()
|
| D | hlo_parser.cc | 1274 optional<int64> split_dimension; in ParseInstructionRhs() local
|
| D | hlo_instructions.cc | 726 const absl::optional<int64>& split_dimension) in HloAllToAllInstruction()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 2857 XlaOp XlaBuilder::AllToAll(XlaOp operand, int64 split_dimension, in AllToAll()
2872 XlaOp operand, int64 split_dimension, int64 concat_dimension, in AllToAllArray()
2923 XlaOp XlaBuilder::AllToAllTuple(XlaOp operand, int64 split_dimension, in AllToAllTuple()
4677 XlaOp AllToAll(const XlaOp operand, int64 split_dimension, in AllToAll()
4685 XlaOp AllToAllTuple(const XlaOp operand, int64 split_dimension, in AllToAllTuple()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 1649 int64 channel_id, absl::optional<int64> split_dimension) { in GetPerGroupCollectiveOpsCreator()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 827 int split_dimension = sharding_it->first; in CreateOrGetSplitNodesForInputSharding() local
|