code: def test_jax_FunctionGraph_once():
      """Make sure that an output is only computed once when it's referenced multiple times."""
      from pytensor.link.jax.dispatch import jax_funcify
  
      x = vector("x")
      y = vector("y")
  
      class TestOp(Op):
          def __init__(self):
              self.called = 0
  
          def make_node(self, *args):
              return Apply(self, list(args), [x.type() for x in args])
  
          def perform(self, inputs, outputs):
              for i, inp in enumerate(inputs):
                  outputs[i][0] = inp[0]
  
      @jax_funcify.register(TestOp)
      def jax_funcify_TestOp(op, **kwargs):
          def func(*args, op=op):
              op.called += 1
              return list(args)
  
          return func
  
      op1 = TestOp()
      op2 = TestOp()
  
      q, r = op1(x, y)
      outs = op2(q + r, q + r)
  
      out_fg = FunctionGraph([x, y], outs, clone=False)
      assert len(out_fg.outputs) == 2
  
      out_jx = jax_funcify(out_fg)
  
      x_val = np.r_[1, 2].astype(config.floatX)
      y_val = np.r_[2, 3].astype(config.floatX)
  
      res = out_jx(x_val, y_val)
      assert len(res) == 2
      assert op1.called == 1
      assert op2.called == 1
  
      res = out_jx(x_val, y_val)
      assert len(res) == 2
      assert op1.called == 2
      assert op2.called == 2
  
imports: from collections.abc import Iterable
  from functools import partial
  from typing import Callable, Optional
  import numpy
  import pytest
  from pytensor.compile.function import function
  from pytensor.compile.mode import get_mode
  from pytensor.compile.sharedvalue import SharedVariable, shared
  from pytensor.configdefaults import config
  from pytensor.graph.basic import Apply
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import Op, get_test_value
  from pytensor.ifelse import ifelse
  from pytensor.raise_op import assert_op
  from pytensor.tensor.type import dscalar, scalar, vector
label_desc: Use image node_id 3 for calling a global function with example usage: test_jax_FunctionGraph_once() and returns: func, Apply, list