| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/ |
| H A D | group_assignment_test.cc | 60 StatusOr<GroupAssignment> group_assignment = GroupAssignment::FromMLIR( in CreateGroupAssignment() local
75 StatusOr<GroupAssignment> group_assignment = GroupAssignment::FromMLIR( in CreateGroupAssignment() local
132 GroupAssignment group_assignment = in TEST() local
145 GroupAssignment group_assignment = in TEST() local
158 GroupAssignment group_assignment = in TEST() local
175 GroupAssignment group_assignment = in TEST() local
|
| H A D | dtensor_allreduce_sum_optimization.cc | 96 mlir::DenseIntElementsAttr group_assignment; in CheckReduceAndSumOptimizationCriteria() local
148 mlir::Value group_assignment = first_reduction_op.group_assignment(); in OptimizeAllReduceAndSum() local
306 mlir::Operation* group_assignment = in CheckWhileLoopOptimizationCriteria() local
|
| H A D | dtensor_allreduce_combine_optimization.cc | 371 StatusOr<GroupAssignment> group_assignment = GroupAssignment::FromMLIR( in CombineAllReduceOpsOfSameTypeAndGroupAssignment() local
519 mlir::Value group_assignment = all_reduce.group_assignment(); in CombineAllReduceOpsOfSameType() local
533 for (mlir::Value group_assignment : group_assignments) { in CombineAllReduceOpsOfSameType() local
|
| H A D | group_assignment.cc | 166 GroupAssignment group_assignment( in FromMLIR() local
|
| H A D | collectives.cc | 315 auto group_assignment = in EmitAllReduce() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
| H A D | collective_ops.py | 71 def assign_group_v2(group_assignment, device_index, base_key): argument
469 group_assignment=None, argument
499 def all_to_all_v3(communicator, t, group_assignment=None, timeout_seconds=None): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | collective_ops.cc | 488 const Tensor& group_assignment = context->input(0); in Compute() local
525 static Status ComputeGroupKey(const Tensor& group_assignment, in ComputeGroupKey()
1066 const Tensor& group_assignment, in FillCollectiveParams()
1191 Tensor group_assignment = c->input(2); in ComputeAsync() local
1241 Tensor group_assignment = c->input(2); in ComputeAsync() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ |
| H A D | tpu_optimizer.py | 36 group_assignment=None): argument
73 def _verify_and_get_subgroup_size(self, group_assignment, num_shards): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ops/ |
| H A D | tpu_ops.py | 44 group_assignment=None, argument
89 def cross_replica_sum(x, group_assignment=None, name=None): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/ops/ |
| H A D | tpu_cross_replica_ops.cc | 36 ShapeHandle group_assignment = c->input(1); in __anon173a7a2a0102() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| H A D | collective_lowering.cc | 120 const mlir::DenseIntElementsAttr& group_assignment, int32 key_base, in EmitCollectiveReduce()
600 const mlir::DenseIntElementsAttr group_assignment = in LowerAllGatherOp() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | cross_replica_ops.cc | 29 const std::vector<int64_t>& group_assignment, in Convert()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| H A D | legalize_tf_collective.cc | 109 DenseIntElementsAttr group_assignment; in ConvertReplicaGroups() local
|
| H A D | legalize_tf.cc | 6095 DenseIntElementsAttr group_assignment; in matchAndRewrite() local
|