| 74 | } |
| 75 | |
| 76 | void Compute(tensorflow::OpKernelContext* context) override { |
| 77 | const Tensor& method_tensor = context->input(1); |
| 78 | OP_REQUIRES(context, TensorShapeUtils::IsScalar(method_tensor.shape()), |
| 79 | errors::InvalidArgument("`method` should be a scalar string: ", |
| 80 | method_tensor.shape())); |
| 81 | // For now, farmhash64 is the only function supported. |
| 82 | const string& method = method_tensor.scalar<string>()(); |
| 83 | OP_REQUIRES( |
| 84 | context, method == "farmhash64", |
| 85 | errors::InvalidArgument("Unsupported fingerprint method: ", method)); |
| 86 | |
| 87 | const Tensor& input = context->input(0); |
| 88 | OP_REQUIRES( |
| 89 | context, TensorShapeUtils::IsVectorOrHigher(input.shape()), |
| 90 | errors::InvalidArgument("`data` should have at least one dimension: ", |
| 91 | input.shape())); |
| 92 | |
| 93 | const int64 dim0 = input.shape().dim_size(0); |
| 94 | const int64 dim1 = input.shape().num_elements() / dim0; |
| 95 | |
| 96 | Tensor* output; |
| 97 | OP_REQUIRES_OK(context, |
| 98 | context->allocate_output( |
| 99 | 0, TensorShape{dim0, kFingerprintSize}, &output)); |
| 100 | |
| 101 | if (input.dtype() == DT_STRING) { |
| 102 | if (dim1 > 1) { |
| 103 | Tensor temp; |
| 104 | OP_REQUIRES_OK(context, context->allocate_temp( |
| 105 | DT_UINT8, |
| 106 | TensorShape{input.shape().num_elements(), |
| 107 | kFingerprintSize}, |
| 108 | &temp)); |
| 109 | // `temp` is a matrix of shape {input.num_elements, fingerprint_size}, |
| 110 | // and each row contains the fingerprint value of corresponding string. |
| 111 | // To compute fingerprints of multiple strings, this op fingerprints the |
| 112 | // buffer containing the string fingerprints. |
| 113 | FarmhashFingerprint64(input.flat<tstring>(), temp.tensor<uint8, 2>()); |
| 114 | FarmhashFingerprint64(static_cast<const Tensor&>(temp).shaped<uint8, 2>( |
| 115 | {dim0, dim1 * kFingerprintSize}), |
| 116 | output->matrix<uint8>()); |
| 117 | } else { |
| 118 | // In case dim1 == 1, each string computes into its own fingerprint |
| 119 | // value. There is no need to fingerprint twice. |
| 120 | FarmhashFingerprint64(input.flat<tstring>(), output->matrix<uint8>()); |
| 121 | } |
| 122 | } else { |
| 123 | auto data = input.bit_casted_shaped<uint8, 2>( |
| 124 | {dim0, dim1 * DataTypeSize(input.dtype())}); |
| 125 | FarmhashFingerprint64(data, output->matrix<uint8>()); |
| 126 | } |
| 127 | } |
| 128 | |
| 129 | private: |
| 130 | static constexpr int kFingerprintSize = sizeof(uint64); |
nothing calls this directly
no test coverage detected