Searched refs:XlaScopedShardingAssignment (Results 1 – 8 of 8) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | spmd_manual_sharding_ops.cc | 63 xla::XlaScopedShardingAssignment assign_sharding(ctx->builder(), in Compile()
75 xla::XlaScopedShardingAssignment assign_sharding(ctx->builder(), manual); in Compile()
118 xla::XlaScopedShardingAssignment assign_sharding(ctx->builder(), manual); in Compile()
126 xla::XlaScopedShardingAssignment assign_sharding(ctx->builder(), in Compile()
|
| D | sharding_op.cc | 37 xla::XlaScopedShardingAssignment no_sharding(ctx->builder(), in Compile()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_compilation_device.cc | 137 xla::XlaScopedShardingAssignment assign_sharding(b, op_sharding); in Compute()
|
| D | xla_compiler.cc | 240 xla::XlaScopedShardingAssignment assign_sharding(builder, sharding); in BuildComputation()
339 xla::XlaScopedShardingAssignment assign_sharding(builder, sharding); in BuildComputation()
395 xla::XlaScopedShardingAssignment assign_sharding(builder, op_sharding); in BuildComputation()
1072 xla::XlaScopedShardingAssignment assign_tuple_sharding( in BuildArguments()
1083 xla::XlaScopedShardingAssignment assign_sharding( in BuildArguments()
1096 xla::XlaScopedShardingAssignment assign_sharding( in BuildArguments()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/xla/ |
| D | infeed_op.cc | 144 xla::XlaScopedShardingAssignment clear_sharding(b, absl::nullopt); in Compile()
|
| D | host_compute_ops.cc | 137 xla::XlaScopedShardingAssignment assign_sharding(b, device_sharding); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.h | 1480 class XlaScopedShardingAssignment {
1482 XlaScopedShardingAssignment(xla::XlaBuilder* builder, in XlaScopedShardingAssignment() function
1488 XlaScopedShardingAssignment(const XlaScopedShardingAssignment&) = delete;
1489 XlaScopedShardingAssignment& operator=(const XlaScopedShardingAssignment&) =
1492 ~XlaScopedShardingAssignment() { SetSharding(prev_sharding_); } in ~XlaScopedShardingAssignment()
|
| D | xla_builder.cc | 1697 XlaScopedShardingAssignment scoped_sharding(this, sharding); in Infeed()
1715 XlaScopedShardingAssignment scoped_sharding(this, in Infeed()
1811 XlaScopedShardingAssignment scoped_sharding(this, OpSharding()); in Outfeed()
|