| 93 | } |
| 94 | |
| 95 | Status ModelWeights::Initialize(OpKernelContext* const context) { |
| 96 | OpInputList sparse_indices_inputs; |
| 97 | TF_RETURN_IF_ERROR( |
| 98 | context->input_list("sparse_indices", &sparse_indices_inputs)); |
| 99 | OpInputList sparse_weights_inputs; |
| 100 | TF_RETURN_IF_ERROR( |
| 101 | context->input_list("sparse_weights", &sparse_weights_inputs)); |
| 102 | OpInputList dense_weights_inputs; |
| 103 | TF_RETURN_IF_ERROR( |
| 104 | context->input_list("dense_weights", &dense_weights_inputs)); |
| 105 | |
| 106 | OpOutputList sparse_weights_outputs; |
| 107 | TF_RETURN_IF_ERROR(context->output_list("out_delta_sparse_weights", |
| 108 | &sparse_weights_outputs)); |
| 109 | |
| 110 | OpOutputList dense_weights_outputs; |
| 111 | TF_RETURN_IF_ERROR( |
| 112 | context->output_list("out_delta_dense_weights", &dense_weights_outputs)); |
| 113 | |
| 114 | for (int i = 0; i < sparse_weights_inputs.size(); ++i) { |
| 115 | Tensor* delta_t; |
| 116 | TF_RETURN_IF_ERROR(sparse_weights_outputs.allocate( |
| 117 | i, sparse_weights_inputs[i].shape(), &delta_t)); |
| 118 | // Convert the input vector to a row matrix in internal representation. |
| 119 | auto deltas = delta_t->shaped<float, 2>({1, delta_t->NumElements()}); |
| 120 | deltas.setZero(); |
| 121 | sparse_weights_.emplace_back(FeatureWeightsSparseStorage{ |
| 122 | sparse_indices_inputs[i].flat<int64>(), |
| 123 | sparse_weights_inputs[i].shaped<float, 2>( |
| 124 | {1, sparse_weights_inputs[i].NumElements()}), |
| 125 | deltas}); |
| 126 | } |
| 127 | |
| 128 | // Reads in the weights, and allocates and initializes the delta weights. |
| 129 | const auto initialize_weights = |
| 130 | [&](const OpInputList& weight_inputs, OpOutputList* const weight_outputs, |
| 131 | std::vector<FeatureWeightsDenseStorage>* const feature_weights) { |
| 132 | for (int i = 0; i < weight_inputs.size(); ++i) { |
| 133 | Tensor* delta_t; |
| 134 | TF_RETURN_IF_ERROR( |
| 135 | weight_outputs->allocate(i, weight_inputs[i].shape(), &delta_t)); |
| 136 | // Convert the input vector to a row matrix in internal |
| 137 | // representation. |
| 138 | auto deltas = delta_t->shaped<float, 2>({1, delta_t->NumElements()}); |
| 139 | deltas.setZero(); |
| 140 | feature_weights->emplace_back(FeatureWeightsDenseStorage{ |
| 141 | weight_inputs[i].shaped<float, 2>( |
| 142 | {1, weight_inputs[i].NumElements()}), |
| 143 | deltas}); |
| 144 | } |
| 145 | return Status::OK(); |
| 146 | }; |
| 147 | |
| 148 | return initialize_weights(dense_weights_inputs, &dense_weights_outputs, |
| 149 | &dense_weights_); |
| 150 | } |
| 151 | |
| 152 | // Computes the example statistics for given example, and model. Defined here |
nothing calls this directly
no test coverage detected