Home
last modified time | relevance | path

Searched refs:state_variable_vector (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_embedding_load_retrieve_ops.cc123 tpu::kMaxAuxiliaryParameterCount + 1> &state_variable_vector) { in GetStateVariables()
186 state_variable_vector[i].push_back(absl::MakeConstSpan( in GetStateVariables()
200 state_variable_vector[i].push_back(absl::Span<const float>()); in GetStateVariables()
209 tpu::kMaxAuxiliaryParameterCount + 1> state_variable_vector; in Compute() local
211 GetStateVariables(ctx, state_variable_vector); in Compute()
212 const int num_tables = state_variable_vector[0].size(); in Compute()
223 state_variable_vector[i][table_id].size(); in Compute()
225 const_cast<float*>(state_variable_vector[i][table_id].data()); in Compute()
267 tpu::kMaxAuxiliaryParameterCount + 1> &state_variable_vector, in GetStateVariables()
288 for (auto& v : state_variable_vector) { in GetStateVariables()
[all …]
H A Dtpu_embedding_load_retrieve_ops.h51 tpu::kMaxAuxiliaryParameterCount + 1>& state_variable_vector);
84 tpu::kMaxAuxiliaryParameterCount + 1>& state_variable_vector,