tico 0.3.0.dev260930__tar.gz → 0.3.0.dev261001__tar.gz

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tico
-Version: 0.3.0.dev260930
+Version: 0.3.0.dev261001
 Summary: Convert Exported Torch Module To Circle
 License: This file provides full text of licenses used in this project
         

tico-0.3.0.dev261001/tico/_version.py

@@ -0,0 +1 @@
+__version__ = "0.3.0.dev261001"

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/recipes/debug/static_gemma4_runtime.py

@@ -52,6 +52,7 @@ from tico.quantization.wrapq.wrappers.gemma4.static_vision_profile import (
 from tico.quantization.wrapq.wrappers.gemma4.utils import (
     assert_gemma4_e2b_no_moe,
     build_decode_attention_mask,
+    gemma4_layer_type_head_dim,
     StaticGemma4Layout,
 )
 
@@ -413,11 +414,7 @@ class StaticGemma4Runtime:
             is_sliding = layer_type == "sliding_attention"
 
             # Per-layer-type head_dim (HF Gemma4TextAttention.__init__)
-            global_head_dim = getattr(self.text_config, "global_head_dim", None)
-            if not is_sliding and global_head_dim:
-                head_dim = int(global_head_dim)
-            else:
-                head_dim = int(self.text_config.head_dim)
+            head_dim = gemma4_layer_type_head_dim(self.text_config, layer_type)
 
             # Per-layer-type num_kv_heads
             # use_alternative_attention = attention_k_eq_v and not is_sliding

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/recipes/debug/wrapper_smoke/cases/gemma4.py

@@ -346,6 +346,7 @@ def _build_text_config(
         attention_k_eq_v=attention_k_eq_v,
         num_kv_shared_layers=num_kv_shared_layers,
         hidden_size_per_layer_input=hidden_size_per_layer_input,
+        allow_global_per_layer_attribute_access=True,
     )
     return _set_eager_attention(text_cfg)
 
@@ -373,6 +374,7 @@ def _build_vision_config(*, size_profile: str) -> Any:
             use_clipped_linears=False,
             rope_parameters={"rope_type": "default", "rope_theta": 100.0},
             standardize=True,
+            allow_global_per_layer_attribute_access=True,
         )
     elif size_profile in _GEMMA4_E2B_WIDTH_PROFILES:
         vision_cfg = Gemma4VisionConfig(
@@ -393,6 +395,7 @@ def _build_vision_config(*, size_profile: str) -> Any:
             use_clipped_linears=True,
             rope_parameters={"rope_type": "default", "rope_theta": 100.0},
             standardize=False,
+            allow_global_per_layer_attribute_access=True,
         )
     else:
         raise AssertionError(f"Unhandled Gemma4 size profile: {size_profile}")
@@ -409,6 +412,7 @@ def _build_vision_patch_embedder_config(*, size_profile: str) -> Any:
             hidden_size=32,
             patch_size=4,
             position_embedding_size=8,
+            allow_global_per_layer_attribute_access=True,
         )
     return _build_vision_config(size_profile=size_profile)
 

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/recipes/debug/wrapper_smoke/cases/gemma4_assistant.py

@@ -24,6 +24,9 @@ from tico.quantization.recipes.debug.wrapper_smoke.case import (
     ForwardInput,
     WrapperSmokeCase,
 )
+from tico.quantization.wrapq.wrappers.gemma4_assistant.utils import (
+    assistant_layer_type_head_dim,
+)
 
 
 def _has_gemma4_assistant() -> CaseAvailability:
@@ -73,6 +76,7 @@ def make_tiny_gemma4_assistant_config() -> Any:
         num_kv_shared_layers=2,
         hidden_size_per_layer_input=0,
         vocab_size_per_layer_input=0,
+        allow_global_per_layer_attribute_access=True,
     )
     config = Gemma4AssistantConfig(
         text_config=text_cfg,
@@ -105,14 +109,16 @@ def _make_assistant_sample(model: torch.nn.Module, kv_len: int) -> ForwardInput:
     """Create one synthetic draft-one assistant sample."""
     text_cfg = model.config.get_text_config()
     kv_heads = int(text_cfg.num_key_value_heads)
+    full_head_dim = assistant_layer_type_head_dim(text_cfg, "full_attention")
+    sliding_head_dim = assistant_layer_type_head_dim(text_cfg, "sliding_attention")
     shared_kv_states = {
         "full_attention": (
-            torch.randn(1, kv_heads, kv_len, int(text_cfg.global_head_dim)),
-            torch.randn(1, kv_heads, kv_len, int(text_cfg.global_head_dim)),
+            torch.randn(1, kv_heads, kv_len, full_head_dim),
+            torch.randn(1, kv_heads, kv_len, full_head_dim),
         ),
         "sliding_attention": (
-            torch.randn(1, kv_heads, kv_len, int(text_cfg.head_dim)),
-            torch.randn(1, kv_heads, kv_len, int(text_cfg.head_dim)),
+            torch.randn(1, kv_heads, kv_len, sliding_head_dim),
+            torch.randn(1, kv_heads, kv_len, sliding_head_dim),
         ),
     }
     return ForwardInput(

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/recipes/export/gemma4_assistant.py

@@ -300,9 +300,9 @@ def write_gemma4_assistant_manifest(
             "num_hidden_layers": int(text_config.num_hidden_layers),
             "num_attention_heads": int(text_config.num_attention_heads),
             "num_key_value_heads": kv_heads,
-            "head_dim": int(text_config.head_dim),
-            "global_head_dim": int(
-                getattr(text_config, "global_head_dim", 0) or text_config.head_dim
+            "head_dim": assistant_layer_type_head_dim(text_config, "sliding_attention"),
+            "global_head_dim": assistant_layer_type_head_dim(
+                text_config, "full_attention"
             ),
             "layer_types": list(text_config.layer_types),
             "sliding_window": int(getattr(text_config, "sliding_window", 0) or 0),

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_for_causal_lm.py

@@ -81,6 +81,7 @@ def _make_text_config() -> Gemma4TextConfig:
         enable_moe_block=False,
         # Enable logit softcapping to exercise that code path.
         final_logit_softcapping=30.0,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_for_conditional_generation.py

@@ -79,6 +79,7 @@ def _make_vision_config() -> Gemma4VisionConfig:
         use_clipped_linears=False,
         rope_parameters={"rope_type": "default", "rope_theta": 100.0},
         standardize=True,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")
@@ -110,6 +111,7 @@ def _make_text_config() -> Gemma4TextConfig:
         enable_moe_block=False,
         # Enable logit softcapping to exercise that code path.
         final_logit_softcapping=30.0,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_model.py

@@ -88,6 +88,7 @@ def _make_vision_config() -> Gemma4VisionConfig:
         use_clipped_linears=False,
         rope_parameters={"rope_type": "default", "rope_theta": 100.0},
         standardize=True,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")
@@ -117,6 +118,7 @@ def _make_text_config() -> Gemma4TextConfig:
         attention_dropout=0.0,
         use_cache=False,
         enable_moe_block=False,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_multimodal_embedder.py

@@ -105,6 +105,7 @@ def main():
         rms_norm_eps=1e-6,
         use_clipped_linears=False,
         rope_parameters={"rope_type": "default", "rope_theta": 100.0},
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(vision_cfg, "_attn_implementation"):
         setattr(vision_cfg, "_attn_implementation", "eager")
@@ -130,6 +131,7 @@ def main():
         attention_dropout=0.0,
         use_cache=False,
         enable_moe_block=False,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(text_cfg, "_attn_implementation"):
         setattr(text_cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_vision_model.py

@@ -116,6 +116,7 @@ def main():
         use_clipped_linears=False,
         rope_parameters={"rope_type": "default", "rope_theta": 100.0},
         standardize=True,
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_vision_patch_embedder.py

@@ -125,6 +125,7 @@ def main():
         rms_norm_eps=1e-6,
         use_clipped_linears=False,
         rope_parameters={"rope_type": "default", "rope_theta": 100.0},
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/examples/gemma4/quantize_vision_pooler.py

@@ -117,6 +117,7 @@ def main():
         rms_norm_eps=1e-6,
         use_clipped_linears=False,
         rope_parameters={"rope_type": "default", "rope_theta": 100.0},
+        allow_global_per_layer_attribute_access=True,
     )
     if not hasattr(cfg, "_attn_implementation"):
         setattr(cfg, "_attn_implementation", "eager")

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/wrappers/gemma4/utils.py

@@ -443,6 +443,26 @@ def extract_text_config(config: Any) -> Any:
     return getattr(config, "text_config", config)
 
 
+def gemma4_layer_type_head_dim(text_config: Any, layer_type: str) -> int:
+    """Return the attention head dim used by one Gemma4 text layer type.
+
+    Mirrors ``Gemma4TextAttention.__init__``. transformers 5.15+ consumes the
+    ``global_head_dim`` constructor argument into a full-attention override in
+    ``per_layer_config`` and no longer exposes it as a plain attribute, so the
+    per-layer view is authoritative whenever the config owns a layer of the
+    requested type. Older configs expose ``global_head_dim`` directly; sliding
+    layers always use ``head_dim``.
+    """
+    per_layer_config = getattr(text_config, "per_layer_config", None)
+    layer_types = tuple(getattr(text_config, "layer_types", ()) or ())
+    if per_layer_config is not None and layer_type in layer_types:
+        return int(per_layer_config[layer_type].head_dim)
+    global_head_dim = getattr(text_config, "global_head_dim", None)
+    if layer_type == "full_attention" and global_head_dim:
+        return int(global_head_dim)
+    return int(text_config.head_dim)
+
+
 def ensure_static_shape(
     name: str, tensor: torch.Tensor, expected: Iterable[int]
 ) -> None:

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/quantization/wrapq/wrappers/gemma4_assistant/utils.py

@@ -19,7 +19,10 @@ from typing import Any, Mapping
 import torch
 import torch.nn as nn
 
-from tico.quantization.wrapq.wrappers.gemma4.utils import assert_gemma4_e2b_no_moe
+from tico.quantization.wrapq.wrappers.gemma4.utils import (
+    assert_gemma4_e2b_no_moe,
+    gemma4_layer_type_head_dim,
+)
 
 
 HF_GEMMA4_ASSISTANT_CLASS_PATH = (
@@ -45,15 +48,13 @@ def extract_assistant_text_config(model_or_config: Any) -> Any:
 def assistant_layer_type_head_dim(text_config: Any, layer_type: str) -> int:
     """Return the attention head dim used by one assistant layer type.
 
-    Mirrors ``Gemma4TextAttention.__init__``: full-attention layers use
-    ``global_head_dim`` when it is set, sliding layers always use ``head_dim``.
+    Rejects layer types outside the supported assistant contract and then
+    delegates to :func:`gemma4_layer_type_head_dim`, which resolves the
+    per-layer ``head_dim`` across transformers versions.
     """
     if layer_type not in SUPPORTED_ASSISTANT_LAYER_TYPES:
         raise ValueError(f"Unsupported Gemma4 assistant layer type: {layer_type!r}.")
-    global_head_dim = getattr(text_config, "global_head_dim", None)
-    if layer_type == "full_attention" and global_head_dim:
-        return int(global_head_dim)
-    return int(text_config.head_dim)
+    return gemma4_layer_type_head_dim(text_config, layer_type)
 
 
 def assistant_shared_kv_num_heads(text_config: Any) -> int:

tico-0.3.0.dev261001/tico/serialize/operators/op_bitwise_and.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# 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
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.circle_mapping import extract_torch_dtype
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.errors import NotYetSupportedError
+from tico.utils.validate_args_kwargs import BitwiseAndArgs
+
+
+@register_node_visitor
+class BitwiseAndVisitor(NodeVisitor):
+    """
+    Lower `aten.bitwise_and.Tensor` on boolean tensors to Circle LOGICAL_AND.
+
+    For boolean operands `bitwise_and` is identical to `logical_and`, which is
+    what `tensor_a & tensor_b` produces (e.g. the attention-mask construction in
+    recent `transformers` releases). Integer bitwise AND has no Circle operator
+    and is rejected.
+    """
+
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.bitwise_and.Tensor]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = BitwiseAndArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        other = args.other
+
+        for operand in (input, other):
+            operand_dtype = extract_torch_dtype(operand)
+            if operand_dtype != torch.bool:
+                raise NotYetSupportedError(
+                    "aten.bitwise_and.Tensor is only supported for bool operands "
+                    f"(got {operand_dtype})"
+                )
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.LOGICAL_AND,
+            self._op_codes,
+        )
+
+        inputs = [input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.LogicalAndOptions
+        )
+        option = circle.LogicalAndOptions.LogicalAndOptionsT()
+
+        operator.builtinOptions = option
+
+        return operator

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/utils/pytree_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 from typing import Any, Dict, Tuple
 
 import torch
@@ -70,6 +71,54 @@ __all__ = [
 ##################################################################################
 
 
+# ---------------------------------------------------------------------------
+# aux_data (pytree context) serialization
+# ---------------------------------------------------------------------------
+#
+# `torch.export.save` serializes the TreeSpec of every pytree node to JSON.
+# Our aux_data carries `torch.dtype`, `torch.device`, and class objects, which
+# json cannot dump, so each registration passes these converters as
+# `to_dumpable_context` / `from_dumpable_context`.
+
+_DTYPE_TAG = "__torch_dtype__"
+_DEVICE_TAG = "__torch_device__"
+_CLASS_TAG = "__class__"
+
+
+def _to_dumpable_value(value: Any) -> Any:
+    if isinstance(value, torch.dtype):
+        # str(torch.float32) == "torch.float32"
+        return {_DTYPE_TAG: str(value).removeprefix("torch.")}
+    if isinstance(value, torch.device):
+        return {_DEVICE_TAG: str(value)}
+    if isinstance(value, type):
+        return {_CLASS_TAG: f"{value.__module__}.{value.__qualname__}"}
+    return value
+
+
+def _from_dumpable_value(value: Any) -> Any:
+    if isinstance(value, dict) and len(value) == 1:
+        if _DTYPE_TAG in value:
+            return getattr(torch, value[_DTYPE_TAG])
+        if _DEVICE_TAG in value:
+            return torch.device(value[_DEVICE_TAG])
+        if _CLASS_TAG in value:
+            module_name, _, qualname = value[_CLASS_TAG].rpartition(".")
+            obj: Any = importlib.import_module(module_name)
+            for attr in qualname.split("."):
+                obj = getattr(obj, attr)
+            return obj
+    return value
+
+
+def _aux_to_dumpable(aux_data: Dict[str, Any]) -> Dict[str, Any]:
+    return {key: _to_dumpable_value(value) for key, value in aux_data.items()}
+
+
+def _aux_from_dumpable(aux_data: Dict[str, Any]) -> Dict[str, Any]:
+    return {key: _from_dumpable_value(value) for key, value in aux_data.items()}
+
+
 # ---------------------------------------------------------------------------
 # StaticCache
 # ---------------------------------------------------------------------------
@@ -124,6 +173,8 @@ def register_static_cache():
             _flatten_static_cache,
             _unflatten_static_cache,
             serialized_type_name=f"{StaticCache.__module__}.{StaticCache.__name__}",
+            to_dumpable_context=_aux_to_dumpable,
+            from_dumpable_context=_aux_from_dumpable,
             flatten_with_keys_fn=_flatten_with_keys_static_cache,
         )
         fx_pytree.register_pytree_flatten_spec(
@@ -226,6 +277,8 @@ def register_static_layer():
             _flatten_static_layer,
             _unflatten_static_layer,
             serialized_type_name=f"{StaticLayer.__module__}.{StaticLayer.__name__}",
+            to_dumpable_context=_aux_to_dumpable,
+            from_dumpable_context=_aux_from_dumpable,
             flatten_with_keys_fn=_flatten_with_keys_static_layer,
         )
         fx_pytree.register_pytree_flatten_spec(
@@ -299,6 +352,8 @@ def register_dynamic_layer():
             _flatten_dynamic_layer,
             _unflatten_dynamic_layer,
             serialized_type_name=f"{DynamicLayer.__module__}.{DynamicLayer.__name__}",
+            to_dumpable_context=_aux_to_dumpable,
+            from_dumpable_context=_aux_from_dumpable,
             flatten_with_keys_fn=_flatten_with_keys_dynamic_layer,
         )
         fx_pytree.register_pytree_flatten_spec(
@@ -435,6 +490,8 @@ def register_dynamic_cache():
             _flatten_dynamic_cache,
             _unflatten_dynamic_cache,
             serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+            to_dumpable_context=_aux_to_dumpable,
+            from_dumpable_context=_aux_from_dumpable,
             flatten_with_keys_fn=_flatten_with_keys_dynamic_cache,
         )
         fx_pytree.register_pytree_flatten_spec(

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/utils/utils.py

@@ -153,6 +153,13 @@ def enforce_type(callable):
         parameters = dict(zip(spec.args, args))
         parameters.update(kwargs)
 
+        # Resolve string annotations to actual types
+        try:
+            type_hints = typing.get_type_hints(callable)
+        except Exception:
+            # Fallback if get_type_hints fails (e.g., missing imports in module)
+            type_hints = spec.annotations or {}
+
         # Return tuple of flattened types.
         # Q) What is flatten?
         # A) Optional/Union is not included. Below are included.
@@ -211,10 +218,10 @@ def enforce_type(callable):
                 continue
 
             assert (
-                name in spec.annotations
+                name in type_hints
             ), f"All parameter require type hints. {name} needs a type hint"
 
-            type_hint = spec.annotations[name]
+            type_hint = type_hints[name]
             type_hint = _flatten_type(type_hint)
             type_check_result = _check_type(value, type_hint)
             if not type_check_result:

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico/utils/validate_args_kwargs.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from dataclasses import dataclass, field
 from typing import List, Optional, TYPE_CHECKING, Union
 
@@ -149,6 +151,17 @@ class AdaptiveAvgPool2dArgs:
     output_size: List[int]
 
 
+@enforce_type
+@dataclass
+class BitwiseAndArgs:
+    """
+    bitwise_and.Tensor(Tensor self, Tensor other) -> Tensor
+    """
+
+    input: torch.fx.Node
+    other: torch.fx.Node
+
+
 @enforce_type
 @dataclass
 class BmmArgs:

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tico
-Version: 0.3.0.dev260930
+Version: 0.3.0.dev261001
 Summary: Convert Exported Torch Module To Circle
 License: This file provides full text of licenses used in this project
         

{tico-0.3.0.dev260930 → tico-0.3.0.dev261001}/tico.egg-info/SOURCES.txt

@@ -508,6 +508,7 @@ tico/serialize/operators/op_arange_start_step.py
 tico/serialize/operators/op_argmax.py
 tico/serialize/operators/op_attention.py
 tico/serialize/operators/op_avg_pool2d.py
+tico/serialize/operators/op_bitwise_and.py
 tico/serialize/operators/op_bmm.py
 tico/serialize/operators/op_cat.py
 tico/serialize/operators/op_circle_gather_nd.py

tico-0.3.0.dev260930/tico/_version.py

@@ -1 +0,0 @@
-__version__ = "0.3.0.dev260930"