| /external/tensorflow/tensorflow/core/ops/ |
| D | tpu_cross_replica_ops.cc | 42 int split_dimension; in __anonb72bce3d0102() local
51 TF_RETURN_IF_ERROR(c->GetAttr("split_dimension", &split_dimension)); in __anonb72bce3d0102()
52 if (split_dimension < 0 || split_dimension >= rank) { in __anonb72bce3d0102()
53 return errors::InvalidArgument("split_dimension ", split_dimension, in __anonb72bce3d0102()
63 in_idx = split_dimension; in __anonb72bce3d0102()
64 } else if (i == split_dimension) { in __anonb72bce3d0102()
|
| D | array_ops.cc | 568 DimensionHandle split_dimension; in __anon7c94107b0a02() local
571 0, c->Rank(input), &split_dimension)); in __anon7c94107b0a02()
574 if (!c->ValueKnown(split_dimension)) { in __anon7c94107b0a02()
581 int64 split_dim = c->Value(split_dimension); in __anon7c94107b0a02()
604 DimensionHandle split_dimension; in __anon7c94107b0b02() local
607 2, c->Rank(input), &split_dimension)); in __anon7c94107b0b02()
621 } else if (size_splits == nullptr && c->ValueKnown(split_dimension)) { in __anon7c94107b0b02()
627 c->Value(split_dimension), in __anon7c94107b0b02()
631 } else if (size_splits == nullptr && !c->ValueKnown(split_dimension)) { in __anon7c94107b0b02()
642 int64 split_dim = c->Value(split_dimension); in __anon7c94107b0b02()
|
| D | ops.pbtxt | 699 name: "split_dimension"
|
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 92 def split(cls, tensor, split_dimension, num_devices): argument
108 if (shape[split_dimension] is not None and
109 shape[split_dimension] < num_devices):
112 (shape, split_dimension, num_devices))
115 tile_assignment_dims[split_dimension] = num_devices
184 def split(tensor, split_dimension, num_devices): argument
185 Sharding.split(tensor, split_dimension, num_devices).apply_to_tensor(tensor)
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 46 split_dimension, argument
72 split_dimension=split_dimension,
87 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/compiler/xla/service/ |
| D | shape_inference.cc | 1992 const Shape& shape, int64 split_dimension, int64 concat_dimension, in InferAllToAllShape() argument
1995 if (split_dimension >= shape.rank() || split_dimension < 0) { in InferAllToAllShape()
1998 split_dimension, ShapeUtil::HumanString(shape)); in InferAllToAllShape()
2005 if (shape.dimensions(split_dimension) % split_count != 0) { in InferAllToAllShape()
2009 shape.dimensions(split_dimension), split_count); in InferAllToAllShape()
2013 new_dimensions[split_dimension] /= split_count; in InferAllToAllShape()
|
| D | shape_inference.h | 134 int64 split_dimension,
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | local_computation_builder.cc | 499 const LocalOp& operand, int64 split_dimension, int64 concat_dimension, in AllToAll() argument
506 return xla::AllToAll(operand.op(), split_dimension, concat_dimension, in AllToAll()
|
| D | local_computation_builder.h | 255 LocalOp AllToAll(const LocalOp& operand, int64 split_dimension,
|
| D | xla_client.py | 1207 split_dimension, argument
1238 return self._client.AllToAll(operand, split_dimension, concat_dimension,
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 2159 XlaOp XlaBuilder::AllToAll(const XlaOp& operand, int64 split_dimension, in AllToAll() argument
2173 ShapeInference::InferAllToAllShape(operand_shape, split_dimension, in AllToAll()
2181 operand_shape.dimensions(split_dimension) / split_count; in AllToAll()
2185 /*stride=*/1, /*dimno=*/split_dimension)); in AllToAll()
3244 XlaOp AllToAll(const XlaOp& operand, int64 split_dimension, in AllToAll() argument
3247 return operand.builder()->AllToAll(operand, split_dimension, concat_dimension, in AllToAll()
|
| D | xla_builder.h | 474 XlaOp AllToAll(const XlaOp& operand, int64 split_dimension,
890 friend XlaOp AllToAll(const XlaOp& operand, int64 split_dimension,
1593 XlaOp AllToAll(const XlaOp& operand, int64 split_dimension,
|
| /external/tensorflow/tensorflow/compiler/xla/g3doc/ |
| D | operation_semantics.md | 58 - The dimension size of the operand on the split_dimension is divisible by
62 <b> `AllToAll(operand, split_dimension, concat_dimension, split_count,
69 | `split_dimension` | `int64` | A value in the interval `[0, |
92 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…
|
| /external/tensorflow/tensorflow/go/op/ |
| D | wrappers.go | 25654 …input tf.Output, group_assignment tf.Output, concat_dimension int64, split_dimension int64, split_…
25658 …interface{}{"concat_dimension": concat_dimension, "split_dimension": split_dimension, "split_count…
|
| /external/tensorflow/tensorflow/core/ops/compat/ |
| D | ops_history.v2.pbtxt | 1578 name: "split_dimension"
1631 name: "split_dimension"
|
| D | ops_history.v1.pbtxt | 1578 name: "split_dimension"
1631 name: "split_dimension"
|