code: def get_model():
      # Make a simple convnet with batch normalization and dropout.
      inputs = keras.Input(shape=(28, 28, 1))
      x = keras.layers.Rescaling(1.0 / 255.0)(inputs)
      x = keras.layers.Conv2D(
          filters=12, kernel_size=3, padding="same", use_bias=False
      )(x)
      x = keras.layers.BatchNormalization(scale=False, center=True)(x)
      x = keras.layers.ReLU()(x)
      x = keras.layers.Conv2D(
          filters=24,
          kernel_size=6,
          use_bias=False,
          strides=2,
      )(x)
      x = keras.layers.BatchNormalization(scale=False, center=True)(x)
      x = keras.layers.ReLU()(x)
      x = keras.layers.Conv2D(
          filters=32,
          kernel_size=6,
          padding="same",
          strides=2,
          name="large_k",
      )(x)
      x = keras.layers.BatchNormalization(scale=False, center=True)(x)
      x = keras.layers.ReLU()(x)
      x = keras.layers.GlobalAveragePooling2D()(x)
      x = keras.layers.Dense(256, activation="relu")(x)
      x = keras.layers.Dropout(0.5)(x)
      outputs = keras.layers.Dense(10)(x)
      model = keras.Model(inputs, outputs)
      return model
  
file_path: keras/guides/distributed_training_with_jax.py
imports: import os
  import jax
  import numpy
  import tensorflow
  import keras
  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: get_model() and returns: model