feat: add broken example with lower

examples/00_translate.py

@@ -1,3 +1,11 @@
+"""
+This example demonstrates how to define a mathematical function using CasADi
+and translate it into JAX-compatible code using the jaxadi library.
+It shows the function's signature in CasADi and its equivalent representation
+in JAX, enabling the use of JAX's automatic differentiation and optimization
+capabilities for the same function.
+"""
+
 import casadi as cs
 
 from jaxadi import translate
@@ -7,9 +15,9 @@
 y = cs.SX.sym("y", 1)
 casadi_function = cs.Function("myfunc", [x, y], [x**2 + y**2 - 1])
 
-print("Signature of the casadi function:")
+print("Signature of the CasADi function:")
 print(casadi_function)
 
-print("Transcribed code:")
+print("Translated JAX function:")
 for cg_str in translate(casadi_function):
     print(cg_str)

examples/01_lower.py

@@ -0,0 +1,17 @@
+import casadi as cs
+
+from jaxadi import lower, translate, declare
+
+# define input variables for the function
+x = cs.SX.sym("x", 10, 10)
+y = cs.SX.sym("y", 10, 10)
+casadi_fn = cs.Function("myfunc", [x, y], [x @ y])
+
+print("Signature of the CasADi function:")
+print(casadi_fn)
+
+# define jax function from casadi one
+jax_fn = declare(translate(casadi_fn))
+
+print("Lowered JAX function:")
+print(lower(jax_fn, casadi_fn))

jaxadi/__init__.py

@@ -1,3 +1,4 @@
 from ._compile import lower
 from ._convert import convert
 from ._translate import translate
+from ._declare import declare

jaxadi/_translate.py

@@ -22,15 +22,8 @@ def translate(func: Function) -> list[str]:
     const_instr = [func.instruction_constant(i) for i in range(n_instr)]
 
     # generate string with complete code
-    codegen = textwrap.dedent(
-        """
-    import jax
-    import jax.numpy as jnp
-
-
-    """
-    )
-    codegen += "@jax.jit\n"
+    codegen = ""
+    # codegen += "@jax.jit\n"
     codegen += f"def evaluate_{func.name()}(*args):\n"
     codegen += "    inputs = args\n"  # combine all inputs into a single list
     codegen += f"    outputs = [jnp.zeros(out) for out in {out_shapes}]\n"  # output variables
@@ -58,4 +51,4 @@ def translate(func: Function) -> list[str]:
     # footer
     codegen += "\n    return outputs\n"
 
-    return codegen.split("\n")
+    return codegen