| /external/tensorflow/tensorflow/core/ops/ |
| D | tpu_cross_replica_ops.cc | 42 int split_dimension; in __anon481f6f7f0102() local
54 TF_RETURN_IF_ERROR(c->GetAttr("split_dimension", &split_dimension)); in __anon481f6f7f0102()
55 if (split_dimension < 0 || split_dimension >= rank) { in __anon481f6f7f0102()
56 return errors::InvalidArgument("split_dimension ", split_dimension, in __anon481f6f7f0102()
68 if (i == split_dimension) { in __anon481f6f7f0102()
|
| D | array_ops.cc | 574 DimensionHandle split_dimension; in __anonf6523ebd0a02() local
577 0, c->Rank(input), &split_dimension)); in __anonf6523ebd0a02()
580 if (!c->ValueKnown(split_dimension)) { in __anonf6523ebd0a02()
587 int64 split_dim = c->Value(split_dimension); in __anonf6523ebd0a02()
610 DimensionHandle split_dimension; in __anonf6523ebd0b02() local
613 2, c->Rank(input), &split_dimension)); in __anonf6523ebd0b02()
627 } else if (size_splits == nullptr && c->ValueKnown(split_dimension)) { in __anonf6523ebd0b02()
633 c->Value(split_dimension), in __anonf6523ebd0b02()
637 } else if (size_splits == nullptr && !c->ValueKnown(split_dimension)) { in __anonf6523ebd0b02()
648 int64 split_dim = c->Value(split_dimension); in __anonf6523ebd0b02()
|
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 93 def split(cls, tensor, split_dimension, num_devices): argument
109 if (shape[split_dimension] is not None and
110 shape[split_dimension] < num_devices):
113 (shape, split_dimension, num_devices))
116 tile_assignment_dims[split_dimension] = num_devices
221 split_dimension, argument
236 Sharding.split(tensor, split_dimension, num_devices).apply_to_tensor(
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 47 split_dimension, argument
73 split_dimension=split_dimension,
88 split_dimension=op.get_attr("concat_dimension"),
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_AllToAll.pbtxt | 37 name: "split_dimension"
52 `split_dimension` and send to the other replicas given group_assignment. After
62 split_dimension=1
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | AllToAll.pbtxt | 30 name: "split_dimension"
83 name: "split_dimension"
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | shape_inference.h | 134 int64 split_dimension,
|
| D | hlo_verifier.cc | 230 if (all_to_all->split_dimension()) { in HandleAllToAll()
253 if (all_to_all->split_dimension()) { in HandleAllToAll()
257 hlo->operand(0)->shape(), *all_to_all->split_dimension(), in HandleAllToAll()
258 *all_to_all->split_dimension(), split_count)); in HandleAllToAll()
|
| D | shape_inference.cc | 2018 const Shape& shape, int64 split_dimension, int64 concat_dimension, in InferAllToAllShape() argument
2021 if (split_dimension >= shape.rank() || split_dimension < 0) { in InferAllToAllShape()
2024 split_dimension, ShapeUtil::HumanString(shape)); in InferAllToAllShape()
2031 if (shape.dimensions(split_dimension) % split_count != 0) { in InferAllToAllShape()
2035 shape.dimensions(split_dimension), split_count); in InferAllToAllShape()
2039 new_dimensions[split_dimension] /= split_count; in InferAllToAllShape()
|
| D | hlo_instructions.cc | 612 const absl::optional<int64>& split_dimension) in HloAllToAllInstruction() argument
615 split_dimension_(split_dimension) {} in HloAllToAllInstruction()
622 shape, new_operands, replica_groups(), channel_id(), split_dimension()); in CloneWithNewOperandsImpl()
649 split_dimension_ == casted_other.split_dimension(); in IdenticalSlowPath()
|
| D | hlo_instruction.cc | 412 absl::optional<int64> split_dimension; in CreateFromProto() local
419 split_dimension = proto.dimensions(0); in CreateFromProto()
426 /*channel_id=*/channel_id, split_dimension); in CreateFromProto()
924 const absl::optional<int64>& split_dimension) { in CreateAllToAll() argument
926 shape, operands, replica_groups, channel_id, split_dimension); in CreateAllToAll()
|
| D | hlo_instructions.h | 386 const absl::optional<int64>& split_dimension);
395 absl::optional<int64> split_dimension() const { return split_dimension_; } in split_dimension() function
|
| D | hlo_parser.cc | 891 optional<int64> split_dimension; in ParseInstructionRhs() local
893 split_dimension = dimensions->at(0); in ParseInstructionRhs()
896 shape, operands, replica_groups, channel_id, split_dimension)); in ParseInstructionRhs()
|
| D | hlo_instruction.h | 645 const absl::optional<int64>& split_dimension = absl::nullopt);
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla_client.py | 1028 split_dimension, argument
1053 return ops.AllToAll(operand, split_dimension, concat_dimension, split_count,
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/tests/ |
| D | ops.mlir | 19 split_dimension = 1 : i64,
32 split_dimension = 1 : i64,
45 split_dimension = 1 : i64,
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 2189 XlaOp XlaBuilder::AllToAll(XlaOp operand, int64 split_dimension, in AllToAll() argument
2203 ShapeInference::InferAllToAllShape(*operand_shape, split_dimension, in AllToAll()
2211 operand_shape->dimensions(split_dimension) / split_count; in AllToAll()
2215 /*stride=*/1, /*dimno=*/split_dimension)); in AllToAll()
3314 XlaOp AllToAll(const XlaOp operand, int64 split_dimension, in AllToAll() argument
3317 return operand.builder()->AllToAll(operand, split_dimension, concat_dimension, in AllToAll()
|
| D | xla_builder.h | 520 XlaOp AllToAll(XlaOp operand, int64 split_dimension, int64 concat_dimension,
928 friend XlaOp AllToAll(XlaOp operand, int64 split_dimension,
1681 XlaOp AllToAll(XlaOp operand, int64 split_dimension, int64 concat_dimension,
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 652 split_dimension=0,
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ir/ |
| D | hlo_ops.cc | 336 auto split_dim_size = type.getDimSize(op.split_dimension().getSExtValue()); in Verify()
|
| D | hlo_ops.td | 541 I64Attr:$split_dimension,
|
| D | hlo_ops_base.td | 697 - The dimension size of the operand on the split_dimension is divisible by
|
| /external/tensorflow/tensorflow/compiler/xla/g3doc/ |
| D | operation_semantics.md | 101 - The dimension size of the operand on the `split_dimension` is divisible by
105 <b> `AllToAll(operand, split_dimension, concat_dimension, split_count,
112 | `split_dimension` | `int64` | A value in the interval `[0, |
135 AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.raw_ops.pbtxt | 85 …argspec: "args=[\'input\', \'group_assignment\', \'concat_dimension\', \'split_dimension\', \'spli…
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.raw_ops.pbtxt | 85 …argspec: "args=[\'input\', \'group_assignment\', \'concat_dimension\', \'split_dimension\', \'spli…
|