From 8efb0fca52579f9241631916149913e6d1a2aab4 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Thu, 27 Jun 2024 22:14:37 +0200 Subject: [PATCH] Doc fixes --- equinox/debug/_dce.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/equinox/debug/_dce.py b/equinox/debug/_dce.py index f13fb451..66091006 100644 --- a/equinox/debug/_dce.py +++ b/equinox/debug/_dce.py @@ -25,7 +25,7 @@ def _register_alive_impl(i, x): def store_dce(x: PyTree, name: Hashable = None): """Used to check whether a PyTree is DCE'd. (That is, whether this code has been - removed in the compiler, due to dead code eliminitation.) + removed in the compiler, due to dead code elimination.) `store_dce` must be used within a JIT'd function, and acts as the identity function. When the JIT'd function is called, then whether each array got DCE'd or @@ -38,11 +38,11 @@ def store_dce(x: PyTree, name: Hashable = None): ```python @jax.jit def f(x): - a, _ = eqxi.store_dce((x**2, x + 1)) + a, _ = eqx.debug.store_dce((x**2, x + 1)) return a f(1) - eqxi.inspect_dce() + eqx.debug.inspect_dce() # Found 1 call to `equinox.debug.store_dce`. # Entry 0: # (i32[], )