class_docstr: CNN layer.
class_docstr_tok: ["CNN","layer","."]
code: def __call__(self, inputs: JTensor) -> JTensor:
      batch_size = inputs.shape[0]
      outputs = inputs
      outputs = self.conv1(outputs)
      outputs = self.pooling1(outputs)[0]
      outputs = self.conv2(outputs)
      outputs = self.pooling2(outputs)[0]
      outputs = jnp.reshape(outputs, [batch_size, -1])
      for _, dense_layer in enumerate(self.dense):
          outputs = dense_layer(outputs)
      return outputs
  
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 2 for calling the CNN obj's underlying member method code with example usage: obj.__call__(inputs) and returns: outputs