[WASM] Add OneHot op #3409
[WASM] Add OneHot op #3409
Conversation
annxingyuan
left a comment
annxingyuan
left a comment
There was a problem hiding this comment.
Thanks @pvaneck ! Just a few minor things.
| } | ||
|
|
||
| registerKernel({ | ||
| kernelName: 'OneHot', |
There was a problem hiding this comment.
Import OneHot from core as well and use here.
tfjs-core/src/ops/one_hot.ts
Outdated
| const res = ENGINE.runKernelFunc( | ||
| forward, inputs as unknown as NamedTensorMap, null /* grad */, OneHot, | ||
| attrs as unknown as NamedAttrMap); | ||
| return reshape(res, outShape); |
There was a problem hiding this comment.
We would like things like reshape to happen within the kernel as well. Could you add logic to reshape in the WASM OneHot typescript implementation?
There was a problem hiding this comment.
I modified the one hot op in core a bit so that flatten is moved into the forward function. To get the correct output shape in the WASM op, I need to pass in the original indices tensor (i.e. not flattened). I hope this is fine.
|
@annxingyuan Thanks for the review! Please take another look when you have a chance. |
annxingyuan
left a comment
annxingyuan
left a comment
There was a problem hiding this comment.
Awesome. Thanks @pvaneck !
3 participants
This adds the OneHot operation to the WebAssembly backend.
Closes: #3114
This change is