code: def jax_funcify_CumOp(op, **kwargs):
      axis = op.axis
      mode = op.mode
  
      def cumop(x, axis=axis, mode=mode):
          if mode == "add":
              return jnp.cumsum(x, axis=axis)
          else:
              return jnp.cumprod(x, axis=axis)
  
      return cumop
  
file_path: pytensor/pytensor/link/jax/dispatch/extra_ops.py
imports: import warnings
  import jax.numpy
  from pytensor.link.jax.dispatch.basic import jax_funcify
  from pytensor.tensor.extra_ops import Bartlett, CumOp, FillDiagonal, FillDiagonalOffset, RavelMultiIndex, Repeat, Unique, UnravelIndex
label_desc: Use image node_id 2 for calling a global function with example usage: jax_funcify_CumOp(op) and returns: cumop, jnp, jnp