fix: outputs

Author: mattephi

jaxadi/_ops.py

@@ -95,5 +95,5 @@
     OP_ATAN2: "jnp.arctan2({0}, {1})",
     OP_CONST: "{0:.16f}",
     OP_INPUT: "inputs[{0}][{1}, {2}]",
-    OP_OUTPUT: "work[{0}][0]",
+    OP_OUTPUT: "{0}[0]",
 }