| 93 | OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!")); |
| 94 | } |
| 95 | void Compute(OpKernelContext *context) override |
| 96 | { |
| 97 | |
| 98 | const Tensor &grad = context->input(0); |
| 99 | auto grad_flat = grad.flat<float>(); |
| 100 | |
| 101 | const Tensor &input = context->input(1); |
| 102 | auto input_flat = input.flat<float>(); |
| 103 | |
| 104 | const Tensor &mu = context->input(2); |
| 105 | auto mu_flat = mu.flat<float>(); |
| 106 | |
| 107 | const Tensor &sigma = context->input(3); |
| 108 | auto sigma_flat = sigma.flat<float>(); |
| 109 | |
| 110 | const TensorShape &input_shape = input.shape(); |
| 111 | |
| 112 | Tensor *output = NULL; |
| 113 | OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &output)); |
| 114 | auto output_flat = output->flat<float>(); |
| 115 | int N = input_flat.size() / times_; |
| 116 | |
| 117 | unsigned long int eid_ = (eid_high_ << 32) + eid_low_; |
| 118 | typedef void (*function)(unsigned long int eid, float *input, float *grad, float *mu, float *sigma, int N, float *output); |
| 119 | dlerror(); |
| 120 | function logprob_grad_kernel = (function)dlsym(lib, "logprob_grad"); |
| 121 | const char *dlsym_error = dlerror(); |
| 122 | OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of logprob_grad failed: ", dlsym_error)); |
| 123 | logprob_grad_kernel(eid_, (float *)input_flat.data(), (float *)grad_flat.data(), (float *)mu_flat.data(), (float *)sigma_flat.data(), N, (float *)output_flat.data()); |
| 124 | }; |
| 125 | |
| 126 | private: |
| 127 | void *lib; |
nothing calls this directly
no outgoing calls
no test coverage detected