code: def make_backbone():
      return keras.Sequential(
          [
              keras.layers.Rescaling(
                  1.0 / 255.0
              ),  # input images are in the range [0, 255]
              keras.layers.Conv2D(
                  filters=12,
                  kernel_size=3,
                  padding="same",
                  use_bias=False,
              ),
              keras.layers.BatchNormalization(scale=False, center=True),
              keras.layers.Activation("relu"),
              keras.layers.Conv2D(
                  filters=24,
                  kernel_size=6,
                  padding="same",
                  use_bias=False,
                  strides=2,
              ),
              keras.layers.BatchNormalization(scale=False, center=True),
              keras.layers.Activation("relu"),
              keras.layers.Conv2D(
                  filters=32,
                  kernel_size=6,
                  padding="same",
                  use_bias=False,
                  strides=2,
                  name="large_k",
              ),
              keras.layers.BatchNormalization(scale=False, center=True),
              keras.layers.Activation("relu"),
          ],
          name="backbone",
      )
  
file_path: keras/examples/demo_jax_distributed.py
imports: import os, pprint, collections
  import jax
  import jax.numpy
  import tensorflow
  import keras
  import numpy
  from tqdm import tqdm
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh
  from jax.sharding import NamedSharding
  from jax.sharding import PartitionSpec
label_desc: Use image node_id 1 for calling a global function with example usage: make_backbone() and returns: keras