Home
last modified time | relevance | path

Searched refs:state_variable_specs (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_embedding_load_retrieve_ops.cc146 std::vector<tpu::StateVariableSpecification> state_variable_specs; in GetStateVariables() local
148 table_descriptor.optimization_parameters(), &state_variable_specs); in GetStateVariables()
162 for (int i = 1; i < state_variable_specs.size(); ++i) { in GetStateVariables()
176 for (int i = 0; i < state_variable_specs.size(); ++i) { in GetStateVariables()
192 for (int i = state_variable_specs.size(); in GetStateVariables()
297 std::vector<tpu::StateVariableSpecification> state_variable_specs; in GetStateVariables() local
299 table_descriptor.optimization_parameters(), &state_variable_specs); in GetStateVariables()
305 num_state_variables[table_id] = state_variable_specs.size(); in GetStateVariables()
307 for (int i = 0; i < state_variable_specs.size(); ++i) { in GetStateVariables()
318 for (int i = state_variable_specs.size(); in GetStateVariables()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/ops/
H A Dtpu_embedding_ops.cc264 std::vector<tpu::StateVariableSpecification> state_variable_specs; in __anoncc2b8c080702() local
267 &state_variable_specs); in __anoncc2b8c080702()
270 for (int i = 1; i < state_variable_specs.size(); ++i) { in __anoncc2b8c080702()
289 for (int i = state_variable_specs.size(); in __anoncc2b8c080702()
335 std::vector<tpu::StateVariableSpecification> state_variable_specs; in __anoncc2b8c080802() local
338 &state_variable_specs)); in __anoncc2b8c080802()
341 (i < state_variable_specs.size() in __anoncc2b8c080802()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
H A Dcombine_tpu_embedding_load_retrieve_pass.cc298 std::vector<tpu::StateVariableSpecification> state_variable_specs; in CombinePerTableParametersForLoad() local
300 opt_params, &state_variable_specs); in CombinePerTableParametersForLoad()
306 CHECK_LE(state_variable_specs.size(), // Crash OK in CombinePerTableParametersForLoad()
316 TF_RET_CHECK(e->dst_input() < state_variable_specs.size()); in CombinePerTableParametersForLoad()
317 TF_RET_CHECK(state_variable_specs[e->dst_input()].has_user_defined() || in CombinePerTableParametersForLoad()
329 if (parameter_num < state_variable_specs.size() && in CombinePerTableParametersForLoad()
330 (state_variable_specs[parameter_num].has_user_defined() || in CombinePerTableParametersForLoad()