class_docstr: MobileNetV1 model.
class_docstr_tok: ["MobileNetV1","model","."]
code: def __call__(self, inputs: jax.Array, is_training: bool) -> jax.Array:
      initial_conv = hk.Conv2D(
          output_channels=32,
          kernel_shape=(3, 3),
          stride=2,
          padding="VALID",
          with_bias=self.with_bias,
      )
  
      out = initial_conv(inputs)
      if self.use_bn:
          bn = hk.BatchNorm(
              create_scale=True, create_offset=True, decay_rate=0.999
          )
          out = bn(out, is_training)
      out = jax.nn.relu(out)
      for i in range(len(self.strides)):
          block = MobileNetV1Block(
              self.channels[i], self.strides[i], self.use_bn
          )
          out = block(out, is_training)
      out = jnp.mean(out, axis=(1, 2))
      out = hk.Flatten()(out)
      out = hk.Linear(self.num_classes, name="logits")(out)
      return out
  
file_path: dm-haiku/haiku/_src/nets/mobilenetv1.py
imports: from collections.abc import Sequence
  from typing import Optional
  from haiku._src import basic
  from haiku._src import batch_norm
  from haiku._src import conv
  from haiku._src import depthwise_conv
  from haiku._src import module
  from haiku._src import reshape
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the MobileNetV1 obj's underlying member method code with example usage: obj.__call__(inputs, is_training) and returns: out