code: def test_invalid_shape_assertion(self):
      arg_i1 = np.bool_(True)
      arg_i32 = np.int32(2)
      res = arg_i32
  
      # This test runs as part of two targets, with and without
      # disabling shape_assertions.
      disabled_shape_assertions_check = (
          "--tf_xla_call_module_disabled_checks=shape_assertions"
          in os.getenv("TF_XLA_FLAGS", "")
      )
      if disabled_shape_assertions_check:
          self.skipTest(
              "Test is N/A when shape_assertions are disabled"
          )
  
      subtest_count = 1
  
      def one_subtest(error_msg: str, module_str: str):
          def f(*args):
              module, version = serialize(module_str)
              return xla.call_module(
                  list(args),
                  version=version,
                  module=module,
                  Tout=[res.dtype],
                  Sout=[res.shape],
                  platforms=[self.testing_platform()],
              )
  
          nonlocal subtest_count
          subtest_count += 1
          with self.subTest(count=subtest_count, error_msg=error_msg):
              with self.assertRaisesRegex(
                  errors.InvalidArgumentError, error_msg
              ):
                  self._assertOpOutputMatchesExpected(
                      f, (arg_i1, arg_i32), (res,)
                  )
  
      one_subtest(
          "expects assert_what .* to be a constant of type tensor<i1>",
          """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg_i1: tensor<i1>, %arg_i32: tensor<i32>) -> tensor<i32> {
      %ok = stablehlo.constant dense<0> : tensor<i32>
      stablehlo.custom_call @shape_assertion(%ok) {
        error_message = "Some error",
        has_side_effect = true
      } : (tensor<i32>) -> ()
      return %arg_i32 : tensor<i32>
    }
  }
  """,
      )
  
      one_subtest(
          "expects static assert_what",
          """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg_i1: tensor<i1>, %arg_i32: tensor<i32>) -> tensor<i32> {
      stablehlo.custom_call @shape_assertion(%arg_i1) {
        error_message = "Some error",
        has_side_effect = true
      } : (tensor<i1>) -> ()
      return %arg_i32 : tensor<i32>
    }
  }
  """,
      )
  
      one_subtest(
          "expects has_side_effect=true",
          """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg_i1: tensor<i1>, %arg_i32: tensor<i32>) -> tensor<i32> {
      %ok = stablehlo.constant dense<false> : tensor<i1>
      stablehlo.custom_call @shape_assertion(%ok) {
        error_message = "Some error",
        has_side_effect = false
      } : (tensor<i1>) -> ()
      return %arg_i32 : tensor<i32>
    }
  }
  """,
      )
  
      one_subtest(
          "expects error_message .* Found specifier {0}",
          """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg_i1: tensor<i1>, %arg_i32: tensor<i32>) -> tensor<i32> {
      %ok = stablehlo.constant dense<false> : tensor<i1>
      stablehlo.custom_call @shape_assertion(%ok) {
        error_message = "Some error {0}",
        has_side_effect = true
      } : (tensor<i1>) -> ()
      return %arg_i32 : tensor<i32>
    }
  }
  """,
      )
  
      one_subtest(
          "expects static error_message_input",
          """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg_i1: tensor<i1>, %arg_i32: tensor<i32>) -> tensor<i32> {
      %ok = stablehlo.constant dense<false> : tensor<i1>
      stablehlo.custom_call @shape_assertion(%ok, %arg_i32) {
        error_message = "Some error {0}",
        has_side_effect = true
      } : (tensor<i1>, tensor<i32>) -> ()
      return %arg_i32 : tensor<i32>
    }
  }
  """,
      )
  
      one_subtest(
          "expects error_message_input .* to be a constant of type tensor<i32>",
          """
  module @jit_fun.1 attributes {jax.uses_shape_polymorphism = true} {
    func.func public @main(%arg_i1: tensor<i1>, %arg_i32: tensor<i32>) -> tensor<i32> {
      %ok = stablehlo.constant dense<false> : tensor<i1>
      %c = stablehlo.constant dense<2.0> : tensor<f32>
      stablehlo.custom_call @shape_assertion(%ok, %c) {
        error_message = "Some error {0}",
        has_side_effect = true
      } : (tensor<i1>, tensor<f32>) -> ()
      return %arg_i32 : tensor<i32>
    }
  }
  """,
      )
  
imports: import os
  import re
  from typing import Optional, Sequence
  import unittest
  from absl.testing import parameterized
  import numpy
  from tensorflow.compiler.mlir.stablehlo import stablehlo
  from tensorflow.compiler.tests import xla_test
  from tensorflow.compiler.tf2xla.ops import gen_xla_ops
  from tensorflow.compiler.tf2xla.python import xla
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import function
  from tensorflow.python.framework import ops
  from tensorflow.python.platform import googletest
  from tensorflow.python.platform import test
label_desc: Use image node_id 27 for calling the XlaCallModuleOpTest obj's underlying member method code with example usage: obj.test_invalid_shape_assertion() and returns: xla