Skip to content

Instantly share code, notes, and snippets.

@dfm

dfm/deterministic_op.ipynb

Created March 2, 2021 23:37
Show Gist options
  • Save dfm/a2db466f46ab931947882b08b2f21558 to your computer and use it in GitHub Desktop.
Save dfm/a2db466f46ab931947882b08b2f21558 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
deterministic_op.ipynb
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@bmorris3
Copy link

bmorris3 commented Jun 25, 2021 •
edited
Loading

Thanks for this! One note I found while experimenting with this on aesara 2.0.12: jax_funcify_JaxOp seems to require an extra keyword arguments node and storage_map, so this tweak makes the code above work for me:

@jax_funcify.register(JaxOp)
def jax_funcify_JaxOp(op, *args, **kwargs):
    func = op.jax_fn
    return func

I hope that's sensible.

@dfm
Copy link
Author

dfm commented Jun 25, 2021

@bmorris3: Yeah - this interface has been a moving target so I haven't been following it too closely, so I'm not sure that I know enough to comment, but seems sensible enough :D

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment