| /external/tensorflow/tensorflow/python/tpu/ |
| D | tpu_optimizer.py | 40 group_assignment=None): argument
70 self._group_assignment = group_assignment
72 def _verify_and_get_subgroup_size(self, group_assignment, num_shards): argument
86 if not group_assignment:
88 if not (isinstance(group_assignment, list) and
89 all(isinstance(i, list) for i in group_assignment)):
91 group_assignment))
94 for g in group_assignment:
101 num_shards, group_assignment))
103 subgroup_size_list = [len(group) for group in group_assignment]
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 41 group_assignment = [list(range(num_shards))]
42 return group_assignment
49 group_assignment=None, argument
67 if group_assignment is None:
68 group_assignment = _create_default_group_assignment()
71 group_assignment,
94 def cross_replica_sum(x, group_assignment=None, name=None): argument
107 if group_assignment is None:
108 group_assignment = _create_default_group_assignment()
110 return gen_tpu_ops.cross_replica_sum(x, group_assignment, name=name)
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_AllToAll.pbtxt | 11 name: "group_assignment"
14 [num_groups, num_replicas_per_group]. `group_assignment[i]` represents the
46 size(group_assignment.get_shape()[1])
52 `split_dimension` and send to the other replicas given group_assignment. After
60 group_assignment=`[[0, 1]]`
|
| D | api_def_CrossReplicaSum.pbtxt | 11 name: "group_assignment"
14 [num_groups, num_replicas_per_group]. `group_assignment[i]` represents the
35 Passing group_assignment=`[[0,2,4,6],[1,3,5,7]]` sets `A, C, E, G` as group 0,
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/ |
| D | CrossReplicaSum.pbtxt | 8 name: "group_assignment"
33 name: "group_assignment"
60 name: "group_assignment"
|
| D | AllToAll.pbtxt | 8 name: "group_assignment"
45 name: "group_assignment"
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | CrossReplicaSum.pbtxt | 8 name: "group_assignment"
33 name: "group_assignment"
|
| D | AllToAll.pbtxt | 8 name: "group_assignment"
45 name: "group_assignment"
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | cross_replica_ops.cc | 29 const std::vector<int64>& group_assignment, in Convert() argument
31 VLOG(1) << "group_assignment size: " << group_assignment.size(); in Convert()
43 int64 replica = group_assignment[num_replica_per_group * g + i]; in Convert()
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.tpu.-cross-shard-optimizer.pbtxt | 21 …argspec: "args=[\'self\', \'opt\', \'reduction\', \'name\', \'group_assignment\'], varargs=None, k…
|
| D | tensorflow.tpu.pbtxt | 33 …argspec: "args=[\'x\', \'group_assignment\', \'name\'], varargs=None, keywords=None, defaults=[\'N…
|
| D | tensorflow.raw_ops.pbtxt | 85 …argspec: "args=[\'input\', \'group_assignment\', \'concat_dimension\', \'split_dimension\', \'spli…
941 …argspec: "args=[\'input\', \'group_assignment\', \'name\'], varargs=None, keywords=None, defaults=…
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_patterns.td | 289 def : Pat<(TF_CrossReplicaSumOp $input, (ConstantLikeMatcher ElementsAttr:$group_assignment)),
291 (CastElementsToI64Elements $group_assignment))>;
297 …lToAllOp AnyRankedTensor:$input, (ConstantLikeMatcher ElementsAttr:$group_assignment), I64Attr:$co…
298 …$split_dimension, $concat_dimension, $split_count, (CastElementsToI64Elements $group_assignment))>;
|
| D | legalize_tf.cc | 5263 DenseIntElementsAttr group_assignment; in matchAndRewrite() local
5264 if (!matchPattern(op.group_assignment(), m_Constant(&group_assignment))) in matchAndRewrite()
5267 hlo::ConvertElementsAttr(group_assignment, rewriter.getIntegerType(64)) in matchAndRewrite()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
| D | hlo_ops_base.td | 1053 Passing group_assignment=`[[0,2,4,6],[1,3,5,7]]` sets `A, C, E, G` as group 0,
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
| D | tf_ops.td | 246 TF_Int32Tensor:$group_assignment,
|
| D | tf_generated_ops.td | 235 `split_dimension` and send to the other replicas given group_assignment. After
243 group_assignment=`[[0, 1]]`
255 [num_groups, num_replicas_per_group]. `group_assignment[i]` represents the
256 replica ids in the ith subgroup.}]>:$group_assignment,
2494 Passing group_assignment=`[[0,2,4,6],[1,3,5,7]]` sets `A, C, E, G` as group 0,
2502 [num_groups, num_replicas_per_group]. `group_assignment[i]` represents the
2503 replica ids in the ith subgroup.}]>:$group_assignment
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.raw_ops.pbtxt | 85 …argspec: "args=[\'input\', \'group_assignment\', \'concat_dimension\', \'split_dimension\', \'spli…
941 …argspec: "args=[\'input\', \'group_assignment\', \'name\'], varargs=None, keywords=None, defaults=…
|
| /external/tensorflow/tensorflow/go/op/ |
| D | wrappers.go | 45556 func CrossReplicaSum(scope *Scope, input tf.Output, group_assignment tf.Output) (output tf.Output) {
45563 input, group_assignment,
51322 func AllToAll(scope *Scope, input tf.Output, group_assignment tf.Output, concat_dimension int64, sp…
51330 input, group_assignment,
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/tests/ |
| D | legalize-tf.mlir | 4881 %group_assignment = "tf.Const" () {
4884 …%result = "tf.AllToAll"(%input, %group_assignment) {T = f32, concat_dimension = 1 : i64, split_cou…
|
| /external/tensorflow/tensorflow/core/ops/compat/ |
| D | ops_history.v2.pbtxt | 1556 name: "group_assignment"
1593 name: "group_assignment"
16581 name: "group_assignment"
16606 name: "group_assignment"
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | ops.pbtxt | 668 name: "group_assignment"
9495 name: "group_assignment"
|