rng = np.random.default_rng(1230)
    random_tensor_gen_fn = lambda shape: rng.uniform(
        low=-1.0, high=1.0, size=shape
    ).astype(np.float32)

    class TwoMatmulModel(module.Module):
      """A model with two matmul ops."""

      @def_function.function
      def matmul(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
        """Performs a matrix multiplication.