| 70 | |
| 71 | template <typename Device, typename Distribution> |
| 72 | Status UpdateVariableAndFill( |
| 73 | OpKernelContext* ctx, Distribution dist, int state_input_idx, |
| 74 | bool read_alg_from_state, Algorithm alg, int64 output_size, |
| 75 | typename Distribution::ResultElementType* output_data) { |
| 76 | Var* var = nullptr; |
| 77 | TF_RETURN_IF_ERROR( |
| 78 | LookupResource(ctx, HandleFromInput(ctx, state_input_idx), &var)); |
| 79 | // Use `ScopedUnlockUnrefVar` here instead of `mutex_lock` and `ScopedUnref` |
| 80 | // because the former supports early releasing which is needed by |
| 81 | // `UpdateVariableAndFill_Philox<CPU>` to avoid holding the lock while |
| 82 | // filling. |
| 83 | ScopedUnlockUnrefVar state_var_guard(var); |
| 84 | Tensor* var_tensor = var->tensor(); |
| 85 | TF_RETURN_IF_ERROR(CheckState(*var_tensor)); |
| 86 | auto var_tensor_flat = var_tensor->flat<StateElementType>(); |
| 87 | int64 alg_tag_skip = 0; |
| 88 | if (read_alg_from_state) { |
| 89 | alg_tag_skip = 1; |
| 90 | if (var_tensor_flat.size() < 1) { |
| 91 | return errors::InvalidArgument("Size of tensor must be at least 1"); |
| 92 | } |
| 93 | alg = var_tensor_flat(0); |
| 94 | } |
| 95 | if (alg == RNG_ALG_PHILOX) { |
| 96 | TF_RETURN_IF_ERROR(CheckPhiloxState(*var_tensor, alg_tag_skip)); |
| 97 | TF_RETURN_IF_ERROR(PrepareToUpdateVariable<Device, StateElementType>( |
| 98 | ctx, var_tensor, var->copy_on_read_mode.load())); |
| 99 | UpdateVariableAndFill_Philox<Device, Distribution>()( |
| 100 | ctx, ctx->eigen_device<Device>(), dist, output_size, alg_tag_skip, |
| 101 | &state_var_guard, var_tensor, output_data); |
| 102 | return Status::OK(); |
| 103 | } else { |
| 104 | return errors::InvalidArgument("Unsupported algorithm id: ", alg); |
| 105 | } |
| 106 | } |
| 107 | |
| 108 | // Preconditon: input(0) is an existing resource. |
| 109 | template <typename Device, class Distribution> |
nothing calls this directly
no test coverage detected