Searched refs:single_dim (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 450 single_dim=-1, argument
472 dim=single_dim,
479 single_dim=-1, argument
503 dim=single_dim,
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ops/ |
| D | xla_ops.cc | 1049 int32 single_dim; in __anonb720852f0f02() local
1050 TF_RETURN_IF_ERROR(c->GetAttr("dim", &single_dim)); in __anonb720852f0f02()
1059 if (single_dim < 0 || single_dim == i) { in __anonb720852f0f02()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 5017 OpSharding GetManualSharding(const OpSharding& original, int64_t single_dim) { in GetManualSharding() argument
5019 if (single_dim < 0 || original.type() != OpSharding::OTHER) { in GetManualSharding()
5027 new_tile_shape.push_back(new_tile_shape[single_dim]); in GetManualSharding()
5028 new_tile_shape[single_dim] = 1; in GetManualSharding()
5037 if (i == single_dim) { in GetManualSharding()
5061 XlaBuilder* builder, XlaOp input, int single_dim, in ConvertSpmdFullToShardShape() argument
5070 if (single_dim >= 0 && i != single_dim) { in ConvertSpmdFullToShardShape()
5095 OpSharding manual = GetManualSharding(manual_sharding, single_dim); in ConvertSpmdFullToShardShape()
5106 XlaBuilder* builder, XlaOp input, const Shape& output_shape, int single_dim, in ConvertSpmdShardToFullShape() argument
5115 OpSharding manual = GetManualSharding(manual_sharding, single_dim); in ConvertSpmdShardToFullShape()
|
| D | xla_builder.h | 2804 xla::XlaBuilder* builder, xla::XlaOp input, int single_dim,
2813 int single_dim, const xla::OpSharding& manual_sharding,
|