code: def jax_simple_model(request: Any) -> ModelAndData:
      import jax
  
      def model(x: Any) -> Any:
          return jax.numpy.mean(x, axis=(1, 2))
  
      bounds = (0, 1)
      fmodel = fbn.JAXModel(model, bounds=bounds)
  
      x, _ = fbn.samples(
          fmodel,
          dataset="cifar10",
          batchsize=8,
          data_format="channels_last",
      )
      x = ep.astensor(x)
      y = fmodel(x).argmax(axis=-1)
      return fmodel, x, y
  
imports: from typing import Optional, Callable, Tuple, Dict, Any, List, NamedTuple
  import functools
  import pytest
  import eagerpy
  import foolbox
  import foolbox
label_desc: Use image node_id 21 for calling a global function with example usage: jax_simple_model(request) and returns: fmodel,  x,  y, jax