code: def MakeZeroState(depth_multiplier=1):
      """Makes zeros of shape like x but removing the length (axis 1)."""
  
      def f(x):  # pylint: disable=invalid-name
          if len(x.shape) != 3:
              raise ValueError(
                  f"Layer input should be a rank 3 tensor representing"
                  f" (batch_size, sequence_length, feature_depth); "
                  f"instead got shape {x.shape}."
              )
          return jnp.zeros(
              (x.shape[0], depth_multiplier * x.shape[-1]),
              dtype=jnp.float32,
          )
  
      return base.Fn("MakeZeroState", f)
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import activation_fns
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
label_desc: Use image node_id 1 for calling a global function with example usage: MakeZeroState(depth_multiplier) and returns: base, jnp