1# Copyright 2021 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Ops for creating TaskEligibilityInfo results.""" 15 16import tensorflow as tf 17 18# Ops implemented in C++ 19from fcp.tensorflow import gen_task_eligibility_info_ops 20 21_task_eligibility_info_ops_so = tf.load_op_library( 22 tf.compat.v1.resource_loader.get_path_to_datafile( 23 "./_task_eligibility_info_ops.so")) 24 25 26def create_task_eligibility_info(version, task_names, task_weights): 27 """Outputs a serialized `TaskEligibilityInfo` proto based on the given inputs. 28 29 This op is used to generate `TaskEligibilityInfo` protos from a model at 30 runtime, since TF Mobile does not support the standard TensorFlow ops for 31 encoding/decoding protos. 32 33 See the `TaskEligibilityInfo` and `TaskWeight` proto message documentation for 34 more information. 35 36 Args: 37 version: an int64 value to place in the `TaskEligibilityInfo.version` field. 38 task_names: a rank-1 string tensor containing the task names to assign 39 weights to. Each entry in this tensor will be combined with the 40 corresponding entry into the `task_weights` tensor at the same index, to 41 form a `TaskWeight` message. 42 task_weights: a rank-1 float tensor containing the task weight for each task 43 (see `task_names`). Note: this tensor must have the same number of 44 elements as `task_names`. 45 46 Returns: 47 a string tensor containing the serialized proto. 48 """ 49 # Convert the inputs to tensors, as a convenience to callers. This ensures 50 # that they can easily pass regular Python or numpy types in addition to 51 # actual tensors. 52 version = tf.convert_to_tensor(version, dtype=tf.int64, name="version") 53 task_names = tf.convert_to_tensor( 54 task_names, dtype=tf.string, name="task_names") 55 task_weights = tf.convert_to_tensor( 56 task_weights, dtype=tf.float32, name="task_weights") 57 return gen_task_eligibility_info_ops.create_task_eligibility_info( 58 version=version, task_names=task_names, task_weights=task_weights) 59