ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/layers/unet/builder.py

@@ -0,0 +1,50 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+# Builder utils for individual components.
+
+import ai_edge_torch.generative.layers.unet.model_config as unet_config
+from torch import nn
+
+
+def build_upsampling(config: unet_config.UpSamplingConfig):
+  if config.mode == unet_config.SamplingType.NEAREST:
+    return nn.UpsamplingNearest2d(scale_factor=config.scale_factor)
+  elif config.mode == unet_config.SamplingType.BILINEAR:
+    return nn.UpsamplingBilinear2d(scale_factor=config.scale_factor)
+  else:
+    raise ValueError("Unsupported upsampling type.")
+
+
+def build_downsampling(config: unet_config.DownSamplingConfig):
+  if config.mode == unet_config.SamplingType.AVERAGE:
+    return nn.AvgPool2d(
+        config.kernel_size, config.stride, padding=config.padding
+    )
+  elif config.mode == unet_config.SamplingType.CONVOLUTION:
+    out_channels = (
+        config.in_channels
+        if config.out_channels is None
+        else config.out_channels
+    )
+    padding = (0, 1, 0, 1) if config.padding == 0 else config.padding
+    return nn.Conv2d(
+        config.in_channels,
+        out_channels=out_channels,
+        kernel_size=config.kernel_size,
+        stride=config.stride,
+        padding=padding,
+    )
+  else:
+    raise ValueError("Unsupported downsampling type.")

ai_edge_torch/generative/layers/unet/model_config.py

@@ -0,0 +1,282 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================
+
+# UNet configuration class.
+import dataclasses
+import enum
+from typing import List, Optional
+
+import ai_edge_torch.generative.layers.model_config as layers_cfg
+
+
+@enum.unique
+class SamplingType(enum.Enum):
+  NEAREST = enum.auto()
+  BILINEAR = enum.auto()
+  AVERAGE = enum.auto()
+  CONVOLUTION = enum.auto()
+
+
+@dataclasses.dataclass
+class UpSamplingConfig:
+  mode: SamplingType
+  scale_factor: float
+
+
+@dataclasses.dataclass
+class DownSamplingConfig:
+  mode: SamplingType
+  in_channels: int
+  kernel_size: int
+  stride: int
+  padding: int
+  out_channels: Optional[int] = None
+
+
+@dataclasses.dataclass
+class ResidualBlock2DConfig:
+  in_channels: int
+  hidden_channels: int
+  out_channels: int
+  hidden_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  # Optional time embedding channels if the residual block takes a time embedding context as input
+  time_embedding_channels: Optional[int] = None
+  residual_out_channels: Optional[int] = None
+
+
+@dataclasses.dataclass
+class AttentionBlock2DConfig:
+  dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  attention_config: layers_cfg.AttentionConfig
+  enable_hlfb: bool = True
+  attention_batch_size: int = 1
+  hidden_dim: Optional[int] = None
+
+
+@dataclasses.dataclass
+class CrossAttentionBlock2DConfig:
+  query_dim: int
+  cross_dim: int
+  hidden_dim: int
+  output_dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  attention_config: layers_cfg.AttentionConfig
+  enable_hlfb: bool = True
+  attention_batch_size: int = 1
+
+
+@dataclasses.dataclass
+class FeedForwardBlock2DConfig:
+  dim: int
+  hidden_dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  use_bias: bool
+
+
+@dataclasses.dataclass
+class TransformerBlock2DConfig:
+  pre_conv_normalization_config: layers_cfg.NormalizationConfig
+  attention_block_config: AttentionBlock2DConfig
+  cross_attention_block_config: CrossAttentionBlock2DConfig
+  feed_forward_block_config: FeedForwardBlock2DConfig
+
+
+@dataclasses.dataclass
+class UpDecoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # The dimension of output channels of previous connected block
+  prev_out_channels: Optional[int] = None
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add upsample operation after residual blocks
+  add_upsample: bool = True
+  # Whether to add a conv2d layer after upsample
+  upsample_conv: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[UpSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+
+
+@dataclasses.dataclass
+class SkipUpDecoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  # The dimension of output channels of previous connected block
+  prev_out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add upsample operation after residual blocks
+  add_upsample: bool = True
+  # Whether to add a conv2d layer after upsample
+  upsample_conv: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[UpSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+  sub_block_channels: Optional[tuple] = None
+  hidden_channels: Optional[int] = None
+
+
+@dataclasses.dataclass
+class DownEncoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Padding for the downsampling convolution.
+  padding: int = 1
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add downsample operation after residual blocks
+  add_downsample: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[DownSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+  hidden_channels: Optional[int] = None
+
+
+@dataclasses.dataclass
+class MidBlock2DConfig:
+  in_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Optional time embedding channels if the residual blocks take a time embedding context as input
+  time_embedding_channels: Optional[int] = None
+  # Optional config of attention blocks interleaved with residual blocks
+  attention_block_config: Optional[AttentionBlock2DConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+
+
+@dataclasses.dataclass
+class AutoEncoderConfig:
+  """Configurations of encoder/decoder in the autoencoder model."""
+
+  # The activation type of encoder/decoder blocks.
+  activation_config: layers_cfg.ActivationConfig
+
+  # The output channels of each block.
+  block_out_channels: List[int]
+
+  # Number of channels in the input image.
+  in_channels: int
+
+  # Number of channels in the output.
+  out_channels: int
+
+  # Number of channels in the latent space.
+  latent_channels: int
+
+  # The component-wise standard deviation of the trained latent space computed using the first batch of the
+  # training set. This is used to scale the latent space to have unit variance when training the diffusion
+  # model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+  # diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+  # / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+  # Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
+  scaling_factor: float
+
+  # The layesr number of each encoder/decoder block.
+  layers_per_block: int
+
+  # The normalization config.
+  normalization_config: layers_cfg.NormalizationConfig
+
+  # The configuration of middle blocks, that is, after the last block of encoder and before the first block of decoder.
+  mid_block_config: MidBlock2DConfig
+
+
+@dataclasses.dataclass
+class DiffusionModelConfig:
+  """Configurations of Diffusion model."""
+
+  # Number of channels in the input tensor.
+  in_channels: int
+
+  # Number of channels in the output tensor.
+  out_channels: int
+
+  # The output channels of each block.
+  block_out_channels: List[int]
+
+  # The layesr number of each block.
+  layers_per_block: int
+
+  # The padding to use for the downsampling.
+  downsample_padding: int
+
+  # Normalization config used in residual blocks.
+  residual_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation config used in residual blocks
+  residual_activation_type: layers_cfg.ActivationType
+
+  # The batch size used in transformer blocks, for attention layers.
+  transformer_batch_size: int
+
+  # The number of attention heads used in transformer blocks.
+  transformer_num_attention_heads: int
+
+  # The dimension of cross attention used in transformer blocks.
+  transformer_cross_attention_dim: int
+
+  # Normalization config used in prev conv layer of transformer blocks.
+  transformer_pre_conv_norm_config: layers_cfg.NormalizationConfig
+
+  # Normalization config used in transformer blocks.
+  transformer_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation type of feed forward used in transformer blocks.
+  transformer_ff_activation_type: layers_cfg.ActivationType
+
+  # Number of layers in mid block.
+  mid_block_layers: int
+
+  # Dimension of time embedding.
+  time_embedding_dim: int
+
+  # Time embedding dimensions for blocks.
+  time_embedding_blocks_dim: int
+
+  # Normalization config used for final layer
+  final_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation type used in final layer
+  final_activation_type: layers_cfg.ActivationType
+
+  # Whether to enable StableHLO composite ops in the model.
+  enable_hlfb: bool = False

ai_edge_torch/generative/quantize/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# 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.
+# ==============================================================================

ai_edge_torch/generative/quantize/example.py

@@ -0,0 +1,47 @@
+# Copyright 2024 The AI Edge Torch Authors. 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.
+# ==============================================================================
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.gemma import gemma1
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.generative.quantize import quant_recipes
+from ai_edge_torch.generative.utilities import model_builder
+import numpy as np
+import torch
+
+
+def main():
+  # Build a PyTorch model as usual
+  config = gemma1.get_fake_model_config()
+  model = model_builder.DecoderOnlyModel(config).eval()
+  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+  tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
+  tokens[0, :4] = idx
+  input_pos = torch.arange(0, 10, dtype=torch.int)
+  kv = kv_utils.KVCache.from_model_config(config)
+
+  # Create a quantization recipe to be applied to the model
+  quant_config = quant_recipes.full_int8_dynamic_recipe()
+  print(quant_config)
+
+  # Convert with quantization
+  edge_model = ai_edge_torch.convert(
+      model, (tokens, input_pos, kv), quant_config=quant_config
+  )
+  edge_model.export("/tmp/gemma_2b_quantized.tflite")
+
+
+if __name__ == "__main__":
+  main()

ai_edge_torch/generative/quantize/quant_attrs.py

@@ -0,0 +1,68 @@
+# Copyright 2024 The AI Edge Torch Authors. 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.
+# ==============================================================================
+
+import enum
+
+
+@enum.unique
+class Dtype(enum.Enum):
+  """Data types and precision of tensors."""
+
+  FP32 = enum.auto()
+  FP16 = enum.auto()
+  INT8 = enum.auto()
+
+
+@enum.unique
+class Algorithm(enum.Enum):
+  """Algorithm used to calculate quantization parameters.
+
+  Attributes:
+    MIN_MAX: Maps the min/max of floating point space to the min/max of
+      quantized space and quantize uniformly.
+    FLOAT_CAST: Casts a float to another float of a different type.
+  """
+
+  MIN_MAX = enum.auto()
+  FLOAT_CAST = enum.auto()
+
+
+@enum.unique
+class Mode(enum.Enum):
+  """Mode of quantization.
+
+  Attributes:
+    DYNAMIC_RANGE: Quantize activations during runtime and weights statically to
+      perform computation in integers.
+    WEIGHT_ONLY: Quantize weights statically and dequantize during runtime to
+      perform computation in floating points.
+  """
+
+  DYNAMIC_RANGE = enum.auto()
+  WEIGHT_ONLY = enum.auto()
+
+
+@enum.unique
+class Granularity(enum.Enum):
+  """Granularity of quantization parameters.
+
+  Attributes:
+    NONE: Granularity not applicable to this quantization scheme.
+    CHANNELWISE: Or per-channel quantization. Each channel of relevant tensors
+      is quantized independently of one another.
+  """
+
+  NONE = enum.auto()
+  CHANNELWISE = enum.auto()

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -0,0 +1,154 @@
+# Copyright 2024 The AI Edge Torch Authors. 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 dataclasses import dataclass
+from typing import Optional, Union
+
+from ai_edge_torch.generative.quantize import quant_attrs
+from ai_edge_torch.generative.quantize import supported_schemes
+
+
+@dataclass
+class LayerQuantRecipe:
+  """Quantization recipe for a single Edge Generative API layer (e.g. Attention).
+
+  Generic layer-scoped quantization recipe that specifies how this layer should
+  be quantized by the Edge Generative API. This is applicable to layers
+  implemented
+  in ai_edge_torch/generative/layers/. Combinations of attributes that are not
+  supported during runtime will be detected when .verify() is called.
+
+  Attributes:
+    activation_dtype: Desired data type of activation tensors.
+    weight_dtype: Desired data type of weight tensors.
+    mode: Type of quantization.
+    algorithm: Algorithm for calculating quantization parameters.
+    granularity: Granularity of quantization.
+  """
+
+  activation_dtype: quant_attrs.Dtype
+  weight_dtype: quant_attrs.Dtype
+  mode: quant_attrs.Mode
+  algorithm: quant_attrs.Algorithm
+  granularity: quant_attrs.Granularity
+
+  def __str__(self):
+    return (
+        f'(a:{self.activation_dtype.name}, '
+        f'w:{self.weight_dtype.name}, '
+        f'{self.mode.name}, '
+        f'{self.algorithm.name}, '
+        f'{self.granularity.name})'
+    )
+
+  __repr__ = __str__
+
+  def verify(self):
+    """Checks if all attributes configured are supported in runtime.
+
+    Raises:
+      ValueError: If any attributes are incompatible.
+    """
+    is_valid = False
+    for supported in supported_schemes.get_supported_layer_schemes():
+      if (
+          self.activation_dtype == supported[0]
+          and self.weight_dtype == supported[1]
+          and self.mode == supported[2]
+          and self.algorithm == supported[3]
+          and self.granularity == supported[4]
+      ):
+        is_valid = True
+        break
+
+    if not is_valid:
+      raise ValueError(
+          'Unsupported LayerQuantRecipe configuration. See'
+          ' get_supported_recipe_matrix()'
+      )
+
+
+@dataclass
+class GenerativeQuantRecipe:
+  """Quantization recipe for a model composed of the Edge Generative API layers.
+
+  Some layers can be specified with different `LayerQuantRecipe` for each block
+  by
+  providing a dictionary keyed by the TransformerBlock index, e.g. attention
+  and feedforward. For example,
+
+  ```
+  default = LayerQuantRecipeA
+  attention = { 2: LayerQuantRecipeB }
+  feedforward = { 3: LayerQuantRecipeC }
+  ```
+
+  will apply LayerQuantRecipeA to the entire model, overriden by
+  LayerQuantRecipeB for the TransformerBlock[2].attention layer and
+  LayerQuantRecipeC for the TransformerBlock[3].feedforward layer. Any config
+  with invalid indices will be ignored.
+
+  Attributes:
+    default: The quantization recipe for global scope of the model.
+    embedding: Recipe for the embedding table.
+    attention: Recipe for the attention blocks. This could be specified with
+      different LayerQuantRecipe for each block by providing a dictionary keyed
+      by the TransformerBlock index.
+    feedforward: Recipe for the feedforward layers. This could be specified with
+      different LayerQuantRecipe for each block by providing a dictionary keyed
+      by the TransformerBlock index.
+  """
+
+  default: Optional[LayerQuantRecipe] = None
+  embedding: Optional[LayerQuantRecipe] = None
+  attention: Union[
+      Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
+  ] = None
+  feedforward: Union[
+      Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
+  ] = None
+
+  def __str__(self):
+    return f"""GenerativeQuantRecipe(
+  Default: {self.default}
+  Embedding: {self.embedding}
+  Attention: {self.attention}
+  Feedforward: {self.feedforward}
+)"""
+
+  __repr__ = __str__
+
+  def verify(self):
+    """Checks if the recipe configured can be supported in runtime.
+
+    Raises:
+      ValueError: If the recipe configured is invalid or unsupported.
+    """
+    if self.default is not None:
+      self.default.verify()
+    if self.embedding is not None:
+      self.embedding.verify()
+    if self.attention is not None:
+      if isinstance(self.attention, dict):
+        for recipe in self.attention.values():
+          recipe.verify()
+      else:
+        self.attention.verify()
+    if self.feedforward is not None:
+      if isinstance(self.feedforward, dict):
+        for recipe in self.feedforward.values():
+          recipe.verify()
+      else:
+        self.feedforward.verify()

ai_edge_torch/generative/quantize/quant_recipe_utils.py

@@ -0,0 +1,62 @@
+# Copyright 2024 The AI Edge Torch Authors. 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.
+# ==============================================================================
+
+"""Helper functions to construct custom quantization recipes.
+
+These are intended for more advanced users who want to configure their own
+quantization recipes. For pre-constructed recipes, use `quant_recipes.py`
+instead.
+
+Typical usage example:
+
+1. Applying a single layer recipe to the entire model
+
+  quant_recipe.GenerativeQuantRecipe(
+    default=quant_recipe_utils.create_layer_quant_int8_dynamic()
+  )
+"""
+
+from ai_edge_torch.generative.quantize import quant_attrs
+from ai_edge_torch.generative.quantize import quant_recipe
+
+
+def create_layer_quant_int8_dynamic() -> quant_recipe.LayerQuantRecipe:
+  return quant_recipe.LayerQuantRecipe(
+      activation_dtype=quant_attrs.Dtype.FP32,
+      weight_dtype=quant_attrs.Dtype.INT8,
+      mode=quant_attrs.Mode.DYNAMIC_RANGE,
+      algorithm=quant_attrs.Algorithm.MIN_MAX,
+      granularity=quant_attrs.Granularity.CHANNELWISE,
+  )
+
+
+def create_layer_quant_int8_weight_only() -> quant_recipe.LayerQuantRecipe:
+  return quant_recipe.LayerQuantRecipe(
+      activation_dtype=quant_attrs.Dtype.FP32,
+      weight_dtype=quant_attrs.Dtype.INT8,
+      mode=quant_attrs.Mode.WEIGHT_ONLY,
+      algorithm=quant_attrs.Algorithm.MIN_MAX,
+      granularity=quant_attrs.Granularity.CHANNELWISE,
+  )
+
+
+def create_layer_quant_fp16() -> quant_recipe.LayerQuantRecipe:
+  return quant_recipe.LayerQuantRecipe(
+      activation_dtype=quant_attrs.Dtype.FP32,
+      weight_dtype=quant_attrs.Dtype.FP16,
+      mode=quant_attrs.Mode.WEIGHT_ONLY,
+      algorithm=quant_attrs.Algorithm.FLOAT_CAST,
+      granularity=quant_attrs.Granularity.NONE,
+  )