maxtext.models.olmo3 module

maxtext.models.olmo3 module#

Copyright 2023-2026 Google LLC

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

maxtext.models.olmo3.get_attention_type(layer_id)[source]#

Get attention type based on layer ID.

class maxtext.models.olmo3.Olmo3DecoderLayer(*args, **kwargs)[source]#

Bases: Module

Transformer decoder layer that attends to the encoder.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

class maxtext.models.olmo3.Olmo3ScannableBlock(*args, **kwargs)[source]#

Bases: Module

A repeatable block of Olmo 3 decoder layers.

This block applies multiple decoder layers sequentially, using the attention pattern defined by OLMO3_ATTENTION_PATTERN. It’s designed to be used with nn.scan for efficient compilation.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

config#

Config, MaxText model config

mesh#

Mesh, JAX device mesh (used for sharding)

num_of_layers#

int, number of decoder layers in the block

quant#

Optional[Quant], quantization config