code: def compute_output_spec(fn, *args, **kwargs):
      with StatelessScope():
          all_input_ktensors = []
          built_in_types = (
              type(None),
              int,
              float,
              str,
              bool,
              complex,
              bytes,
          )
  
          # First, separate symbolic args from other args
          static_args_idx = []
          static_args = []
          maybe_symbolic_args = []
          static_kwargs = {}
          maybe_symbolic_kwargs = {}
          for idx, arg in enumerate(args):
              if isinstance(arg, built_in_types):
                  static_args_idx.append(idx)
                  static_args.append(arg)
              else:
                  maybe_symbolic_args.append(arg)
          for k, v in kwargs.items():
              if isinstance(v, built_in_types):
                  static_kwargs[k] = v
              else:
                  maybe_symbolic_kwargs[k] = v
  
          # Second, identify all ktensors
          def index_all_ktensors(x):
              if isinstance(x, KerasTensor):
                  all_input_ktensors.append(x)
              return x
  
          # Third, find out if there are dynamic shapes
          (
              maybe_symbolic_args,
              maybe_symbolic_kwargs,
          ) = tree.map_structure(
              index_all_ktensors,
              (maybe_symbolic_args, maybe_symbolic_kwargs),
          )
          none_count = 0
          for x in all_input_ktensors:
              for d in x.shape:
                  if d is None:
                      none_count += 1
  
          def convert_keras_tensor_to_jax(x, fill_value=None):
              if isinstance(x, KerasTensor):
                  shape = list(x.shape)
                  if fill_value:
                      for i, e in enumerate(shape):
                          if e is None:
                              shape[i] = fill_value
                  jax_tensor = jax.ShapeDtypeStruct(
                      shape, dtype=x.dtype
                  )
                  return jax_tensor
              if isinstance(x, types.FunctionType):
  
                  def _fn(*args, **kwargs):
                      out = x(*args, **kwargs)
                      out = convert_keras_tensor_to_jax(
                          out, fill_value=fill_value
                      )
                      return out
  
                  return Partial(_fn)
              if isinstance(x, dict):
                  return {
                      k: convert_keras_tensor_to_jax(
                          v, fill_value=fill_value
                      )
                      for k, v in x.items()
                  }
              if isinstance(x, list):
                  return [
                      convert_keras_tensor_to_jax(
                          xi, fill_value=fill_value
                      )
                      for xi in x
                  ]
              return x
  
          def wrapped_fn(*args, **kwargs):
              rec_args = []
              idx_static = 0
              idx_sym = 0
              i = 0
              while idx_static < len(static_args) or idx_sym < len(
                  args
              ):
                  if i in static_args_idx:
                      rec_args.append(static_args[idx_static])
                      idx_static += 1
                  else:
                      rec_args.append(args[idx_sym])
                      idx_sym += 1
  
                  i += 1
              with StatelessScope():
                  return fn(*rec_args, **kwargs, **static_kwargs)
  
          jax_out = None
          if none_count:
              try:
                  ms_args_1, ms_kwargs_1 = tree.map_structure(
                      lambda x: convert_keras_tensor_to_jax(
                          x, fill_value=83
                      ),
                      (maybe_symbolic_args, maybe_symbolic_kwargs),
                  )
                  _, jax_out_1 = jax.make_jaxpr(
                      wrapped_fn, return_shape=True
                  )(*ms_args_1, **ms_kwargs_1)
  
                  ms_args_2, ms_kwargs_2 = tree.map_structure(
                      lambda x: convert_keras_tensor_to_jax(
                          x, fill_value=89
                      ),
                      (maybe_symbolic_args, maybe_symbolic_kwargs),
                  )
                  _, jax_out_2 = jax.make_jaxpr(
                      wrapped_fn, return_shape=True
                  )(*ms_args_2, **ms_kwargs_2)
  
                  flat_out_1 = tree.flatten(jax_out_1)
                  flat_out_2 = tree.flatten(jax_out_2)
  
                  flat_out = []
                  for x1, x2 in zip(flat_out_1, flat_out_2):
                      if isinstance(x1, jax.ShapeDtypeStruct):
                          if not isinstance(x2, jax.ShapeDtypeStruct):
                              raise ValueError(
                                  "Indeterministic output ordering."
                              )
                          shape = list(x1.shape)
                          for i, e in enumerate(x2.shape):
                              if e != shape[i]:
                                  shape[i] = None
                          flat_out.append(
                              jax.ShapeDtypeStruct(
                                  shape, dtype=x1.dtype
                              )
                          )
                      else:
                          flat_out.append(x1)
                  jax_out = pack_sequence_as(jax_out_1, flat_out)
              except Exception as e:
                  if "[JAX RNG]" in str(e):
                      raise e
                  # Errors can happen when the filled dimensions
                  # are not compatible with the function
                  # (or when the function contains a bug).
                  # In such cases we don't want to confuse users
                  # with random filled dimensions and the like,
                  # so we rerun a pass on the dynamic shapes,
                  # which will likely error out when JAX tries to
                  # validate shapes as fully static.
                  # The error message will be much easier to understand.
                  pass
  
          if jax_out is None:
              (
                  maybe_symbolic_args,
                  maybe_symbolic_kwargs,
              ) = tree.map_structure(
                  convert_keras_tensor_to_jax,
                  (maybe_symbolic_args, maybe_symbolic_kwargs),
              )
              _, jax_out = jax.make_jaxpr(
                  wrapped_fn, return_shape=True
              )(*maybe_symbolic_args, **maybe_symbolic_kwargs)
  
          def convert_jax_spec_to_keras_tensor(x):
              if isinstance(x, jax.ShapeDtypeStruct):
                  return KerasTensor(x.shape, x.dtype)
              return x
  
          output_shape = tree.map_structure(
              convert_jax_spec_to_keras_tensor, jax_out
          )
      return output_shape
  
imports: import types
  import jax
  import jax.numpy
  import numpy
  import tree
  from jax.tree_util import Partial
  from keras.backend.common import KerasVariable
  from keras.backend.common import global_state
  from keras.backend.common import standardize_dtype
  from keras.backend.common.keras_tensor import KerasTensor
  from keras.backend.common.stateless_scope import StatelessScope
  from keras.backend.jax import distribution_lib
  from keras.utils.nest import pack_sequence_as
label_desc: Use image node_id 6 for calling a global function with example usage: compute_output_spec(fn) and returns: output_shape, x, x, x, jax_tensor, Partial, unknown, unknown, fn, KerasTensor, out