Helper function for creating a slot variable.
(primary, val, scope, validate_shape, shape, dtype, slot_config)
| 68 | |
| 69 | |
| 70 | def _create_slot_var(primary, val, scope, validate_shape, shape, dtype, slot_config): |
| 71 | """Helper function for creating a slot variable.""" |
| 72 | |
| 73 | # TODO(lukaszkaiser): Consider allowing partitioners to be set in the current |
| 74 | # scope. |
| 75 | current_partitioner = variable_scope.get_variable_scope().partitioner |
| 76 | variable_scope.get_variable_scope().set_partitioner(None) |
| 77 | # When init from val instead of callable initializer, the shape is expected to |
| 78 | # be None, not <unknown> or any fully defined shape. |
| 79 | shape = shape if callable(val) else None |
| 80 | if resource_variable_ops.is_resource_variable(primary): |
| 81 | use_resource = True |
| 82 | elif isinstance(primary, variables.RefVariable): |
| 83 | use_resource = False |
| 84 | else: |
| 85 | use_resource = None |
| 86 | if isinstance(primary, kv_variable_ops.EmbeddingVariable): |
| 87 | if slot_config is None: |
| 88 | slot = variable_scope.get_embedding_variable_internal( |
| 89 | scope, |
| 90 | initializer=val, |
| 91 | trainable=False, |
| 92 | embedding_dim=shape, |
| 93 | key_dtype=primary._invalid_key_type, |
| 94 | validate_shape=validate_shape, |
| 95 | steps_to_live=primary._steps_to_live, |
| 96 | ht_partition_num=primary._ht_partition_num) |
| 97 | _set_init_op_embedding_type_attr(slot, config_pb2.EmbeddingVariableType.MUTABLE) |
| 98 | else: |
| 99 | filter_strategy = None |
| 100 | if primary._filter_freq != 0: |
| 101 | if primary._max_element_size != 0: |
| 102 | filter_strategy = variables.CBFFilter(filter_freq=primary._filter_freq, |
| 103 | max_element_size=primary._max_element_size, |
| 104 | false_positive_probability=primary._false_positive_probability, |
| 105 | counter_type=primary._counter_type) |
| 106 | else: |
| 107 | filter_strategy = variables.CounterFilter(filter_freq=primary._filter_freq) |
| 108 | if slot_config.slot_type is config_pb2.SlotType.EMBEDDING_VARIABLE: |
| 109 | _set_init_op_slot_num_attr(primary, slot_config.slot_num) |
| 110 | primary._slot_num = slot_config.slot_num |
| 111 | emb_index = primary._emb_index |
| 112 | if primary.block_num > 1: |
| 113 | primary = primary._primary |
| 114 | slot = variable_scope.get_embedding_variable_v2_internal( |
| 115 | scope, initializer=val, trainable=False, |
| 116 | embedding_dim=shape, key_dtype=primary._invalid_key_type, |
| 117 | validate_shape=validate_shape, |
| 118 | evconfig=variables.EmbeddingVariableConfig( |
| 119 | steps_to_live=primary._steps_to_live, |
| 120 | handle_name=primary._block_handle_name, |
| 121 | emb_index=emb_index, |
| 122 | block_num=primary.block_num, |
| 123 | slot_index=slot_config.slot_index, |
| 124 | primary=primary._primary, |
| 125 | slot_num=slot_config.slot_num, |
| 126 | storage_type=primary.storage_type, |
| 127 | storage_path=primary._storage_path, |
no test coverage detected