-
Notifications
You must be signed in to change notification settings - Fork 254
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Port FNet #1164
Port FNet #1164
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good! I think we will have to stop pinning the keras-core version before we can do this. At least one more bug to work through there.
x = tf.cast(input, tf.complex64) | ||
mixing_output = tf.math.real(tf.signal.fft2d(x)) | ||
return tf.cast(mixing_output, input.dtype) | ||
x = (input, ops.zeros_like(input)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe make it clearer the real, imaginary split?
real_in, imaginary_in = (input, ops.zeros_like(input))
real_out, imaginary_out = ops.fft2((real_in, imaginary_in))
return real_out
I think we should be all clear to land this now. Trying it out. |
/gcbrun |
/gcbrun |
/gcbrun |
/gcbrun |
I think we can ignore the tf GPU OOM, unrelated. |
Resolves #1139