index, replication_attr_name);
    arg->set_is_same_data_across_replicas(attr != nullptr && attr.getValue());

    // Currently only support first dimension to be bounded dynamic.
    arg->mutable_is_bounded_dynamic_dim()->Add(
        dynamic_arg_idx_set.contains(index));
  }

  return success();
}

// Populates a TPUCompileMetadataProto with result sharding from a
// `tf_device::ClusterFuncOp`.