| /external/tensorflow/tensorflow/core/ops/ |
| D | tpu_cross_replica_ops.cc | 41 int concat_dimension; in __anonc85edbdb0102() local
47 TF_RETURN_IF_ERROR(c->GetAttr("concat_dimension", &concat_dimension)); in __anonc85edbdb0102()
49 if (concat_dimension < 0 || concat_dimension >= rank) { in __anonc85edbdb0102()
50 return errors::InvalidArgument("concat_dimension ", concat_dimension, in __anonc85edbdb0102()
65 if (i == concat_dimension) { in __anonc85edbdb0102()
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_AllToAll.pbtxt | 31 name: "concat_dimension"
54 blocks along `concat_dimension` as the output.
61 concat_dimension=0
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 46 concat_dimension, argument
72 concat_dimension=concat_dimension,
87 concat_dimension=op.get_attr("split_dimension"),
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 102 mlir::LogicalResult CreateConcatOp(const int concat_dimension, in CreateConcatOp() argument
111 mlir::DenseElementsAttr::get(concat_dim_type, concat_dimension); in CreateConcatOp()
123 if (input_type.getShape()[concat_dimension] == in CreateConcatOp()
128 shape[concat_dimension] = shape[concat_dimension] * inputs.size(); in CreateConcatOp()
426 int concat_dimension = sharding.tile_assignment_dimensions_size() - 1; in HandleTileShardedOutputs() local
429 --concat_dimension; in HandleTileShardedOutputs()
439 CreateConcatOp(concat_dimension, location, in HandleTileShardedOutputs()
450 --concat_dimension; in HandleTileShardedOutputs()
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | AllToAll.pbtxt | 26 name: "concat_dimension"
79 name: "concat_dimension"
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/ |
| D | AllToAll.pbtxt | 26 name: "concat_dimension"
79 name: "concat_dimension"
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_constant_folding_test.cc | 113 int concat_dimension; in TEST_F() member
128 dimensions[test_config.concat_dimension] = csize; in TEST_F()
135 dimensions[test_config.concat_dimension] = concat_size; in TEST_F()
138 shape, operands, test_config.concat_dimension)); in TEST_F()
|
| D | shape_inference.h | 143 int64 concat_dimension,
|
| D | shape_inference.cc | 2062 const Shape& shape, int64 split_dimension, int64 concat_dimension, in InferAllToAllShape() argument
2070 if (concat_dimension >= shape.rank() || concat_dimension < 0) { in InferAllToAllShape()
2073 concat_dimension, ShapeUtil::HumanString(shape)); in InferAllToAllShape()
2084 new_dimensions[concat_dimension] *= split_count; in InferAllToAllShape()
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | concat_op_test.cc | 61 int concat_dimension, int dim2, in ConcatHelper() argument
68 concat_dim.scalar<int32>()() = concat_dimension; in ConcatHelper()
159 int concat_dimension, int dim2) { in ConcatManyHelper() argument
166 concat_dim.scalar<int32>()() = concat_dimension; in ConcatManyHelper()
|
| D | quantized_concat_op_test.cc | 252 int concat_dimension, bool same_limits, int dim2) { in ConcatHelper() argument
259 Tensor concat_dim = test::AsScalar<int32>(concat_dimension); in ConcatHelper()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | concat_test.cc | 773 int64 concat_dimension; member
791 ConcatInDim(&builder, {a0, a1}, spec.concat_dimension); in TEST_P()
794 ReferenceUtil::Concat2D(lhs, rhs, spec.concat_dimension); in TEST_P()
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 1326 concat_dimension=0,
1567 concat_dimension=axis,
1580 concat_dimension=0,
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 2858 int64 concat_dimension, int64 split_count, in AllToAll() argument
2864 return AllToAllTuple(operand, split_dimension, concat_dimension, in AllToAll()
2867 return AllToAllArray(operand, split_dimension, concat_dimension, split_count, in AllToAll()
2872 XlaOp operand, int64 split_dimension, int64 concat_dimension, in AllToAllArray() argument
2879 concat_dimension, split_count)); in AllToAllArray()
2896 if (split_dimension == concat_dimension) { in AllToAllArray()
2913 if (i == concat_dimension) { in AllToAllArray()
2924 int64 concat_dimension, int64 split_count, in AllToAllTuple() argument
2939 concat_dimension, split_count) in AllToAllTuple()
2992 return this->ConcatInDim(received, concat_dimension); in AllToAllTuple()
[all …]
|
| D | xla_builder.h | 746 XlaOp AllToAll(XlaOp operand, int64 split_dimension, int64 concat_dimension,
752 int64 concat_dimension, int64 split_count,
1308 int64 concat_dimension, int64 split_count,
1312 int64 concat_dimension, int64 split_count,
1444 int64 concat_dimension, int64 split_count,
2223 XlaOp AllToAll(XlaOp operand, int64 split_dimension, int64 concat_dimension,
2229 int64 concat_dimension, int64 split_count,
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_patterns.td | 297 …put, (ConstantLikeMatcher ElementsAttr:$group_assignment), I64Attr:$concat_dimension, $split_dimen…
298 …(HLO_AllToAllOp $input, $split_dimension, $concat_dimension, $split_count, (CastElementsToI64Eleme…
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/tests/ |
| D | ops.mlir | 20 concat_dimension = 0 : i64,
33 concat_dimension = 0 : i64,
46 concat_dimension = 0 : i64,
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tensor_tracer.py | 1575 x, concat_dimension=0, split_dimension=0,
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
| D | hlo_ops.td | 605 I64Attr:$concat_dimension,
|
| D | hlo_ops_base.td | 810 `concat_dimension`.
|
| /external/tensorflow/tensorflow/compiler/xla/g3doc/ |
| D | operation_semantics.md | 127 `concat_dimension`.
148 <b> `AllToAll(operand, split_dimension, concat_dimension, split_count,
159 | `concat_dimension` | `int64` | a value in the interval `[0, |
178 AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
| D | tf_generated_ops.td | 237 blocks along `concat_dimension` as the output.
244 concat_dimension=0
258 I64Attr:$concat_dimension,
|
| /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 | 51322 func AllToAll(scope *Scope, input tf.Output, group_assignment tf.Output, concat_dimension int64, sp…
51326 …attrs := map[string]interface{}{"concat_dimension": concat_dimension, "split_dimension": split_dim…
|