Method init

src/MaxText/maxengine.py:1535–1588

(abstract_params, page_state)

    # pylint: disable=unused-argument
    def init(abstract_params, page_state):
      x = jnp.ones(
          (int(self.config.per_device_batch_size * self.mesh.size), 1),
          dtype=jnp.int32,
      )
      dummy_image = jnp.ones(
          multimodal_utils.get_dummy_image_shape_for_init(
              self.config.model_name, batch_size=self.config.micro_batch_size_to_train_on
          ),
          dtype=jnp.int32,
      )
      _, cache = self.model.apply(
          abstract_params,
          x,
          x,
          encoder_images=dummy_image if self.config.use_multimodal else None,
          enable_dropout=False,
          model_mode=MODEL_MODE_AUTOREGRESSIVE,
          rngs={"params": rng},
          mutable=["cache"],
          page_state=page_state,
          slot=0,
      )

      next_pos = jnp.zeros(
          (int(self.config.per_device_batch_size * self.mesh.size), 1),
          dtype=jnp.int32,
      )
      generated_tokens = jnp.zeros(
          (int(self.config.per_device_batch_size * self.mesh.size), 1),
          dtype=jnp.int32,
      )
      tokens = jnp.zeros(
          (int(self.config.per_device_batch_size * self.mesh.size), 1),
          dtype=jnp.int32,
      )
      token_logp = jnp.zeros(
          (int(self.config.per_device_batch_size * self.mesh.size), 1),
          dtype=jnp.float32,
      )
      return {
          "logits": jnp.zeros(
              (
                  int(self.config.per_device_batch_size * self.mesh.size),
                  1,
                  self.config.vocab_size,
              )
          ),
          "cache": cache["cache"],
          "next_pos": next_pos,
          "generated_tokens": generated_tokens,
          "tokens": tokens,
          "token_logp": token_logp,
      }

    with nn_partitioning.axis_rules(self.config.logical_axis_rules):
      abstract_outputs = jax.eval_shape(init, self.abstract_params, page_state)