class_docstr: MNIST model based on CNN. xid/62622458.
class_docstr_tok: ["MNIST","model","based","on","CNN",".","xid\/62622458","."]
code: def task(self) -> pax_fiddle.Config[base_task.BaseTask]:
      model_p = pax_fiddle.Config(
          CNNModel,
          network_tpl=pax_fiddle.Config(
              CNN,
              name="mnist_model",
              height=self.HEIGHT,
              width=self.WIDTH,
              num_classes=self.NUM_CLASSES,
              kernel_size=self.KERNEL_SIZE,
          ),
      )
      model_p.ici_mesh_shape = self.ICI_MESH_SHAPE
      model_p.mesh_axis_names = self.MESH_AXIS_NAMES
  
      optimizer_p = pax_fiddle.Config(
          optimizers.ShardedSgd,
          momentum=self.MOMENTUM,
          lr_schedule=pax_fiddle.Config(schedules.Constant, value=1),
      )
      optimizer_p.learning_rate = self.LEARNING_RATE
      learner_p = pax_fiddle.Config(
          learners.Learner,
          loss_name="loss",
          optimizer=optimizer_p,
      )
      task_p = pax_fiddle.Config(
          tasks_lib.SingleTask,
          name="mnist_task",
          model=model_p,
          summary_verbosity=0,
          train=pax_fiddle.Config(
              tasks_lib.SingleTask.Train,
              learner=learner_p,
              num_train_steps=self.TRAIN_STEPS,
              eval_skip_train=True,
          ),
      )
      return task_p
  
imports: from typing import cast
  import jax
  import jax.numpy
  from paxml import base_experiment
  from paxml import base_task
  from paxml import experiment_registry
  from paxml import learners
  from paxml import tasks_lib
  from praxis import base_input
  from praxis import base_layer
  from praxis import base_model
  from praxis import optimizers
  from praxis import pax_fiddle
  from praxis import py_utils
  from praxis import pytypes
  from praxis import schedules
  from praxis.layers import activations
  from praxis.layers import convolutions
  from praxis.layers import linears
  from praxis.layers import poolings
  from pax_examples.dali_pax_input import MnistDaliInput
label_desc: Use image node_id 1 for calling the MnistExperiment obj's underlying member method code with example usage: obj.task() and returns: task_p