commit | ba225e3d375ccca3485b4e3e43a1e72ad6dcaa1d | [log] [tgz] |
---|---|---|
author | Andy Ly <lyandy@google.com> | Fri Nov 06 14:00:03 2020 -0800 |
committer | TensorFlower Gardener <gardener@tensorflow.org> | Fri Nov 06 14:05:11 2020 -0800 |
tree | 2e1a2189fdc24dca00aeb40a837014922cede5d2 | |
parent | 70d22236dbd6bb321b3ec879f3d1f3a2c6990b6c [diff] |
Set _xla_replica_id attribute as a i64 attribute instead of a i32 (NFC). This will match how integer attributes are represented in the TF dialect currently. PiperOrigin-RevId: 341114482 Change-Id: If009fbb7c78ca6911baab601595c851a81555723
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/replicate_to_island.cc b/tensorflow/compiler/mlir/tensorflow/transforms/replicate_to_island.cc index 1b20a27..c051c50 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/replicate_to_island.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/replicate_to_island.cc
@@ -96,7 +96,7 @@ auto result = region.walk([&](Operation* op) -> WalkResult { if (RequiresReplicaIDAttribute(op)) { - op->setAttr(kReplicaIdAttr, builder.getI32IntegerAttr(replica_id)); + op->setAttr(kReplicaIdAttr, builder.getI64IntegerAttr(replica_id)); return WalkResult::advance(); }